semantex-core 1.1.0

//! ColBERT multi-vector encoder for late-interaction retrieval.
//!
//! Uses LateOn-Code-edge via `next-plaid-onnx` for per-token embeddings (48d).
//!
//! # Lazy initialization (E8c)
//!
//! `ColbertEmbedder::new()` is cheap — it only records the model directory.
//! The expensive ONNX session creation is deferred until the first call to
//! `encode_query()` / `encode_documents()`. This shaves ~200-300ms off cold
//! daemon startup; the session is materialized on first query (still well
//! before the user notices, and warmable via [`ColbertEmbedder::warm_up`]).

use anyhow::Result;
use ndarray::Array2;
use next_plaid_onnx::{Colbert, ColbertConfig};
use parking_lot::Mutex;
use std::collections::HashSet;
use std::path::{Path, PathBuf};
use std::sync::OnceLock;
use tokenizers::Tokenizer;

/// Per-token embedding matrix: rows = tokens, cols = embedding dimension (48).
pub type TokenEmbeddings = Array2<f32>;

/// Default per-inference document batch for the single CPU ONNX session.
///
/// Matches `next-plaid-onnx`'s own single-session CPU default and the indexer's
/// outer `PLAID_BATCH`, so one outer encode batch maps to one ORT call. Override
/// with `SEMANTEX_ORT_BATCH` on memory-constrained hosts (smaller = less peak
/// scratch) or to push throughput on large machines (larger = bigger GEMMs).
const DEFAULT_ORT_BATCH: usize = 32;

/// Pure decision logic for the indexing thread-count default, split out from
/// [`index_ort_threads`] so it's testable without touching the real
/// `index::gate::max_concurrent_builds()` (which reads process-global env
/// vars shared with `gate`'s own tests — calling it directly from a parallel
/// test run would race). Identical formula to
/// `embedding::single_vector::default_index_threads`.
fn default_index_threads(cores: usize, max_concurrent_builds: usize) -> usize {
    if max_concurrent_builds <= 1 {
        cores.clamp(2, 8)
    } else {
        (cores / 2).clamp(2, 8)
    }
}

/// ORT intra-op thread count for the **indexing** path (see
/// [`ColbertEmbedder::for_indexing`]). `SEMANTEX_INDEX_ORT_THREADS` overrides;
/// otherwise `clamp(cores / 2, 2, 8)` — enough parallelism to finish a build
/// quickly while leaving cores for the developer's foreground work. Bounded
/// concurrency (the `index::gate` semaphore) keeps total memory in check.
///
/// EXCEPTION: when `index::gate::max_concurrent_builds()` reports this box
/// can only ever run ONE full build at a time (RAM-limited — e.g. a 16 GB /
/// 4-core box, `16 / 16 GB-per-build = 1` slot), there is by construction no
/// sibling build to leave cores for, so the default uses ALL cores instead
/// (still clamped `[2, 8]`). This is the common single-developer/first-index
/// case; bigger boxes where several builds can run concurrently keep the
/// half-the-cores split. Verified live on a 4-core box: without this
/// exception the build used 2 real ORT threads for the whole encode phase
/// (2 idle cores throughout, since only one build could ever be running).
fn index_ort_threads() -> usize {
    let cores = std::thread::available_parallelism().map_or(4, std::num::NonZeroUsize::get);
    let default = default_index_threads(cores, crate::index::gate::max_concurrent_builds());
    crate::config::env_usize("SEMANTEX_INDEX_ORT_THREADS", default)
}

/// Rebuild the exact token-id sequence `next-plaid-onnx` feeds to the model
/// for one document, then drop the same ColBERT-style skiplist ids it filters
/// out of its output embeddings. See the investigation notes on
/// [`ColbertEmbedder::encode_documents_with_ids`] for file:line citations.
///
/// # Note on the truncation branch (`assemble_and_filter_ids`)
///
/// For `LateOn-Code-edge` specifically, `tokenizer.json` embeds its OWN
/// truncation policy (`{"max_length": 2047, "direction": "Right", ...}`, i.e.
/// exactly `document_length - 1`), which the `tokenizers` crate applies
/// automatically inside `.encode()`. So `ids` below can never exceed 2047
/// elements for THIS model, regardless of input length — the
/// `real_len > truncate_limit` branch inside `assemble_and_filter_ids` is
/// unreachable end-to-end via this function for `LateOn-Code-edge` today (it
/// exists for models whose tokenizer does NOT pre-truncate, matching
/// `next-plaid-onnx`'s own defensive logic exactly). It's unit-tested
/// directly on synthetic ids — see `assemble_and_filter_ids`'s test.
pub(crate) fn build_doc_token_ids(alignment: &DocIdAlignment, text: &str) -> Result<Vec<u32>> {
    // Mirrors `next_plaid_onnx::preprocess_texts`: trim, lowercase if the
    // model config requests it.
    let processed = if alignment.do_lower_case {
        text.trim().to_lowercase()
    } else {
        text.trim().to_string()
    };
    let encoding = alignment
        .tokenizer
        .encode(processed.as_str(), true)
        .map_err(|e| anyhow::anyhow!("tokenization error: {e}"))?;
    Ok(assemble_and_filter_ids(
        alignment.document_prefix_id,
        alignment.document_length,
        &alignment.skiplist_ids,
        encoding.get_ids(),
    ))
}

/// Pure id-assembly step, factored out of [`build_doc_token_ids`] so the
/// truncation (`content_prefix_len`/`keep_sep`) branch can be unit-tested
/// directly with a synthetic `ids` sequence, independent of whether any real
/// tokenizer would ever hand back that many raw ids (see the doc comment on
/// [`build_doc_token_ids`] — for `LateOn-Code-edge` it never does).
///
/// `ids` is an already-tokenized `[CLS], content..., [SEP]` sequence (as
/// `tokenizer.encode(text, true).get_ids()` produces). Returns the
/// `[CLS], [D]-marker, content (possibly truncated), optional [SEP]`
/// sequence with ColBERT-style skiplist ids dropped — 1:1 with what
/// `next-plaid-onnx` feeds the model and then filters out of its output rows.
fn assemble_and_filter_ids(
    document_prefix_id: u32,
    document_length: usize,
    skiplist_ids: &HashSet<u32>,
    ids: &[u32],
) -> Vec<u32> {
    // A single (unbatched) encode never pads, so `real_len` is just
    // `ids.len()`; the `.max(1)` guard mirrors next-plaid-onnx's own
    // defensive floor for degenerate (empty) inputs.
    let real_len = ids.len().max(1);
    let truncate_limit = document_length.saturating_sub(1);
    let (content_prefix_len, keep_sep) = if real_len > truncate_limit {
        (truncate_limit.saturating_sub(1), true)
    } else {
        (real_len, false)
    };

    let mut token_ids = Vec::with_capacity(content_prefix_len + 2);
    token_ids.push(ids[0]); // [CLS]
    token_ids.push(document_prefix_id); // [D] marker, inserted at position 1
    token_ids.extend(ids.iter().take(content_prefix_len).skip(1).copied());
    if keep_sep {
        token_ids.push(ids[real_len - 1]); // original [SEP], restored after truncation
    }

    // ColBERT-style punctuation filter: next-plaid-onnx drops any row whose
    // token id is in the skiplist when it flattens the ONNX output, so we
    // drop the matching ids here to stay aligned with its embedding rows.
    token_ids.retain(|id| !skiplist_ids.contains(id));
    token_ids
}

/// ColBERT encoder wrapping LateOn-Code-edge via `next-plaid-onnx`.
///
/// Produces per-token embeddings (N_tokens x 48d) for both queries and documents.
/// The encoder is thread-safe: the inner ONNX session is protected by a
/// `parking_lot::Mutex` (single-threaded inference).
///
/// The ONNX session is created lazily on first encode call (E8c). The
/// initialization itself is serialized by `build_lock` so concurrent first
/// callers do not race to build two ~150 MB ONNX sessions (Finding 10).
pub struct ColbertEmbedder {
    /// Model directory containing `model_int8.onnx` / `tokenizer.json` / `onnx_config.json`.
    model_dir: PathBuf,
    /// Thread count for ORT (read once at construction).
    threads: usize,
    /// CoreML acceleration toggle (read once at construction).
    use_coreml: bool,
    /// Lazily-initialized ONNX session. Constructed on first encode call.
    encoder: OnceLock<Mutex<Colbert>>,
    /// Serializes the build path so only one thread ever calls `build_encoder`.
    /// Readers that observe `encoder.get() == Some(_)` skip this lock entirely
    /// — the fast read-path is preserved.
    build_lock: std::sync::Mutex<()>,
    /// Test/diagnostics: number of times `build_encoder` was actually invoked.
    /// Used by the race-condition test to verify single-build behaviour.
    #[doc(hidden)]
    build_count: std::sync::atomic::AtomicUsize,
    /// Lazily-loaded tokenizer + settings needed to replicate the doc-side
    /// token-id sequence for [`ColbertEmbedder::encode_documents_with_ids`].
    /// Independent `OnceLock` from `encoder` — building this never touches
    /// the ONNX session.
    id_alignment: OnceLock<DocIdAlignment>,
}

/// Tokenizer + config needed to reconstruct, for a single document, the exact
/// token-id sequence `next-plaid-onnx` feeds to the model and then filters
/// down to produce its output embedding rows.
///
/// See the investigation doc comment on
/// [`ColbertEmbedder::encode_documents_with_ids`] for why this is a replica
/// rather than something read directly off the vendored encoder.
///
/// `pub(crate)` (not `pub`): this is an internal replication-of-vendor-
/// behavior detail, not part of the crate's public API. [`StaticTokenEmbedder`]
/// (`embedding::static_token`) is the one other crate-internal consumer —
/// it needs the same alignment `ColbertEmbedder` uses so its own doc-side
/// tokenization can't drift from the contract Task 2/3 established.
pub(crate) struct DocIdAlignment {
    tokenizer: Tokenizer,
    /// Token id for the `[D] ` document marker `next-plaid-onnx` inserts at
    /// position 1 (right after `[CLS]`).
    document_prefix_id: u32,
    /// `onnx_config.json`'s `document_length` — content beyond
    /// `document_length - 1` tokens is truncated, keeping the trailing
    /// `[SEP]`, mirroring `next-plaid-onnx`'s truncation exactly.
    document_length: usize,
    /// Whether the model lowercases text before tokenizing.
    do_lower_case: bool,
    /// Token ids for `onnx_config.json`'s `skiplist_words` (ColBERT-style
    /// punctuation filter) — `next-plaid-onnx` drops these from its output
    /// embedding rows, so this replica must drop the matching ids too.
    skiplist_ids: HashSet<u32>,
}

/// Load the tokenizer + id-alignment settings from `model_dir`'s
/// `tokenizer.json` / `onnx_config.json` files — the same files the ONNX
/// encoder reads.
///
/// Extracted from [`ColbertEmbedder::id_alignment`]'s lazy-init body
/// (behavior-preserving: `id_alignment` now just calls this and caches the
/// result in its own `OnceLock`, exactly as before) so a second, independent
/// caller — [`crate::embedding::static_token::StaticTokenEmbedder`] (Task 5)
/// — can build the identical `DocIdAlignment` once at construction, without
/// needing its own copy of this loading logic or a dependency on
/// `ColbertEmbedder`'s caching.
pub(crate) fn load_doc_id_alignment(model_dir: &Path) -> Result<DocIdAlignment> {
    let tokenizer_path = model_dir.join("tokenizer.json");
    let tokenizer = Tokenizer::from_file(&tokenizer_path).map_err(|e| {
        anyhow::anyhow!("failed to load tokenizer {}: {e}", tokenizer_path.display())
    })?;
    let config = ColbertConfig::from_file(model_dir.join("onnx_config.json"))?;
    // `next-plaid-onnx` itself resolves the `[D]` marker id via its private
    // `ColbertConfig::document_prefix_id` field FIRST, falling back to
    // `tokenizer.token_to_id(document_prefix)` only when that field is
    // `None` (next-plaid-onnx 1.6.2 src/lib.rs:1902-1910). That private
    // field has no public accessor, so we cannot read it from outside the
    // crate — we always take the fallback path here. For the current
    // model (`LateOn-Code-edge`, `document_prefix_id: 50369`) this is
    // verified to resolve to the SAME id as the tokenizer lookup, so
    // there is no active bug today. But the two paths are only known to
    // coincide for this one model: if a future/different model's
    // `onnx_config.json` ever sets a `document_prefix_id` override that
    // does NOT match what the tokenizer resolves for `document_prefix`
    // (e.g. a vocab quirk the override exists specifically to route
    // around), this method would silently report the WRONG id for the
    // marker row. Row-COUNT alignment (`ids.len() == embeddings.nrows()`)
    // would still hold — exactly one marker row is always inserted
    // regardless of its id value — so that check cannot catch this; only
    // the id VALUE would be wrong. If per-token-id distillation (Task 3)
    // ever looks wrong specifically for the `[D]`/marker slot on a new
    // model, this divergence is the first thing to rule out.
    let document_prefix_id = tokenizer
        .token_to_id(&config.document_prefix)
        .ok_or_else(|| {
            anyhow::anyhow!(
                "document prefix token '{}' not found in tokenizer vocabulary",
                config.document_prefix
            )
        })?;
    let skiplist_ids = config
        .skiplist_words
        .iter()
        .filter_map(|word| tokenizer.token_to_id(word))
        .collect();
    Ok(DocIdAlignment {
        tokenizer,
        document_prefix_id,
        document_length: config.document_length,
        do_lower_case: config.do_lower_case,
        skiplist_ids,
    })
}

/// Global singleton for `ColbertEmbedder`.
static GLOBAL_COLBERT: OnceLock<ColbertEmbedder> = OnceLock::new();
/// Serializes initialization of the global ColBERT embedder (double-checked locking).
static COLBERT_INIT_LOCK: parking_lot::Mutex<()> = parking_lot::Mutex::new(());

impl ColbertEmbedder {
    /// Construct a lazy embedder from a model directory containing
    /// `model_int8.onnx` and `tokenizer.json`.
    ///
    /// This call is **cheap** — it only stores the model directory and reads
    /// environment toggles. The actual ONNX session is created on the first
    /// encode call, or via [`ColbertEmbedder::warm_up`] from a background thread.
    ///
    /// Thread count is read from `SEMANTEX_ORT_THREADS` (default 4). This is the
    /// QUERY default: daemons that serve search keep it low so N concurrent
    /// per-project daemons don't multiply ORT intra-op scratch into an OOM. For
    /// the indexing path use [`ColbertEmbedder::for_indexing`] instead, which
    /// picks a higher, machine-adaptive thread count (indexing is throughput-
    /// bound and its concurrency is bounded by the build gate, see `index::gate`).
    /// CoreML acceleration is gated behind `SEMANTEX_COREML=1`.
    ///
    /// # Errors
    ///
    /// Returns an error only if the model directory does not exist. ONNX session
    /// construction errors are deferred to the first encode call.
    pub fn new(model_dir: &Path) -> Result<Self> {
        Self::with_threads(
            model_dir,
            crate::config::env_usize("SEMANTEX_ORT_THREADS", 4),
        )
    }

    /// Construct an embedder tuned for **indexing** (building the PLAID index),
    /// as opposed to serving queries.
    ///
    /// Indexing encodes thousands of chunks and is throughput-bound, so it wants
    /// real CPU parallelism — unlike the query path, which keeps threads low to
    /// bound memory across many concurrent per-project daemons. ORT's intra-op
    /// pool is created per session, independent of any process-global
    /// `OMP_NUM_THREADS` pinned at daemon startup, so an in-process background
    /// build (see `McpServer::spawn_background_index`) gets these threads even
    /// inside a daemon launched in memory-constrained mode.
    ///
    /// Thread count: `SEMANTEX_INDEX_ORT_THREADS` if set, else an adaptive
    /// `clamp(cores / 2, 2, 8)`. Half the cores (capped at 8) keeps a couple
    /// cores free for the developer's foreground work while still finishing a
    /// build in a fraction of the single-threaded time. Memory stays bounded
    /// because only a few full builds run at once (the `index::gate` semaphore).
    pub fn for_indexing(model_dir: &Path) -> Result<Self> {
        Self::with_threads(model_dir, index_ort_threads())
    }

    fn with_threads(model_dir: &Path, threads: usize) -> Result<Self> {
        if !model_dir.exists() {
            anyhow::bail!("ColBERT model dir does not exist: {}", model_dir.display());
        }

        let use_coreml = std::env::var("SEMANTEX_COREML").is_ok_and(|v| v == "1");

        Ok(Self {
            model_dir: model_dir.to_path_buf(),
            threads: threads.max(1),
            use_coreml,
            encoder: OnceLock::new(),
            build_lock: std::sync::Mutex::new(()),
            build_count: std::sync::atomic::AtomicUsize::new(0),
            id_alignment: OnceLock::new(),
        })
    }

    /// Build the inner ONNX session (called lazily on first encode).
    ///
    /// Defaults: int8-quantized model, **explicit CPU execution provider on
    /// every platform**, `SEMANTEX_ORT_THREADS` intra-op threads (default 4),
    /// and a per-inference batch of `SEMANTEX_ORT_BATCH` docs (default
    /// [`DEFAULT_ORT_BATCH`]). Memory profile: session load ~150 MB + per-encode
    /// scratch that scales with batch × sequence length (a few hundred MB at the
    /// default batch). Opt into CoreML on macOS via `SEMANTEX_COREML=1`.
    ///
    /// # Why the explicit CPU provider matters (all platforms)
    ///
    /// `next-plaid-onnx`'s builder defaults to `ExecutionProvider::Auto`, and
    /// its `build()` force-pins intra-op threads to **1** whenever the provider
    /// is `Auto`/`Cuda` for a single session (it assumes a GPU supplies the
    /// parallelism). On Linux/Windows — where no GPU EP is compiled — leaving
    /// the provider at `Auto` therefore silently pinned ColBERT inference to a
    /// **single core**, making CPU indexing 10-30x too slow (a 225-file repo
    /// took >13 min, never finishing). Pinning `Cpu` explicitly keeps our
    /// configured `threads` intra-op count. macOS already pinned a provider, so
    /// it escaped this; the fix simply extends that to every platform.
    ///
    /// # Why batch size 32, not 2
    ///
    /// The previous hardcoded `2` is the library's *parallel-sessions* batch
    /// (small batches suit many 1-thread sessions). For our single multi-thread
    /// CPU session it produced GEMMs too small to engage the intra-op threads.
    /// 32 matches the library's own single-session CPU default and the indexer's
    /// outer `PLAID_BATCH`, so one outer batch maps to one ORT call.
    fn build_encoder(&self) -> Result<Colbert> {
        self.build_count
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);

        let batch_size = crate::config::env_usize("SEMANTEX_ORT_BATCH", DEFAULT_ORT_BATCH);

        #[allow(unused_mut)]
        let mut builder = Colbert::builder(&self.model_dir)
            .with_quantized(true)
            .with_threads(self.threads)
            .with_batch_size(batch_size);

        // Pin an explicit execution provider on EVERY platform — see the method
        // doc comment. On macOS the user may opt into CoreML; everywhere else we
        // force CPU so that a single session keeps its configured intra-op
        // thread count (Auto would silently collapse it to 1).
        #[cfg(target_os = "macos")]
        {
            let provider = if self.use_coreml {
                next_plaid_onnx::ExecutionProvider::CoreML
            } else {
                next_plaid_onnx::ExecutionProvider::Cpu
            };
            builder = builder.with_execution_provider(provider);
        }
        #[cfg(not(target_os = "macos"))]
        {
            let _ = self.use_coreml;
            builder = builder.with_execution_provider(next_plaid_onnx::ExecutionProvider::Cpu);
        }

        builder.build()
    }

    /// Get the ONNX session, initializing it on first call.
    ///
    /// Subsequent calls return the cached session; only the first call pays the
    /// ~200-300ms initialization cost.
    ///
    /// # Concurrency (Finding 10)
    ///
    /// The naive check-then-build pattern previously allowed concurrent first
    /// callers both to observe `encoder.get() == None`, both to call
    /// `build_encoder()` (each allocating a ~150 MB ONNX session), and the
    /// loser's session to be silently dropped — doubling peak RSS and risking
    /// the daemon's RSS guard. We now serialize the build path with a separate
    /// `build_lock` and double-check `OnceLock::get` after acquiring it. The
    /// fast read-path is preserved: once the encoder is set, no lock is taken.
    fn encoder(&self) -> Result<&Mutex<Colbert>> {
        if let Some(enc) = self.encoder.get() {
            return Ok(enc);
        }
        // Serialize the build path. Only one thread ever runs `build_encoder`.
        // Other concurrent first-callers block here until the winner finishes,
        // then re-observe `encoder.get() == Some(_)` and return immediately.
        let _guard = self
            .build_lock
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if let Some(enc) = self.encoder.get() {
            return Ok(enc);
        }
        let built = self.build_encoder()?;
        let _ = self.encoder.set(Mutex::new(built));
        Ok(self
            .encoder
            .get()
            .expect("encoder set in previous statement"))
    }

    /// Force-materialize the ONNX session, optionally running a dummy encode.
    ///
    /// Intended for fire-and-forget background warm threads spawned by the
    /// daemon at startup so that the session is hot before the first user
    /// query arrives. Returns `Ok(())` on success; the caller typically logs
    /// and ignores errors (cold start still works if warm-up fails).
    ///
    /// # Errors
    ///
    /// Returns an error if ONNX session construction or the warm encode fails.
    pub fn warm_up(&self) -> Result<()> {
        let enc = self.encoder()?;
        // Dummy encode to materialize all internal buffers / allocators.
        let _ = enc.lock().encode_queries(&["warmup"])?;
        Ok(())
    }

    /// Get or initialize the global singleton (double-checked locking).
    ///
    /// The first call constructs the embedder (cheap — defers ONNX session);
    /// subsequent calls return the existing instance. Returns an error only
    /// if construction itself fails (e.g., model dir missing).
    pub fn global(model_dir: &Path) -> Result<&'static ColbertEmbedder> {
        if let Some(embedder) = GLOBAL_COLBERT.get() {
            return Ok(embedder);
        }
        let _guard = COLBERT_INIT_LOCK.lock();
        if let Some(embedder) = GLOBAL_COLBERT.get() {
            return Ok(embedder);
        }
        tracing::info!("Initializing global ColBERT encoder singleton (lazy)");
        let embedder = Self::new(model_dir)?;
        let _ = GLOBAL_COLBERT.set(embedder);
        Ok(GLOBAL_COLBERT.get().expect("just set"))
    }

    /// Encode a search query into per-token embeddings `[N_query_tokens, 48]`.
    pub fn encode_query(&self, text: &str) -> Result<TokenEmbeddings> {
        let encoder = self.encoder()?;
        let mut embeddings = encoder.lock().encode_queries(&[text])?;
        embeddings
            .pop()
            .ok_or_else(|| anyhow::anyhow!("encode_queries returned empty result"))
    }

    /// Encode documents into per-token embeddings, one `Array2<f32>` per document.
    pub fn encode_documents(&self, texts: &[String]) -> Result<Vec<TokenEmbeddings>> {
        let encoder = self.encoder()?;
        let refs: Vec<&str> = texts.iter().map(String::as_str).collect();
        let embeddings = encoder.lock().encode_documents(&refs, None)?;
        Ok(embeddings)
    }

    /// Lazily load the tokenizer + id-alignment settings from the same
    /// `tokenizer.json` / `onnx_config.json` files the ONNX encoder reads.
    /// Separate `OnceLock` from `encoder()` — this never builds an ONNX
    /// session, only the (cheap-ish) HF tokenizer.
    fn id_alignment(&self) -> Result<&DocIdAlignment> {
        if let Some(alignment) = self.id_alignment.get() {
            return Ok(alignment);
        }
        let alignment = load_doc_id_alignment(&self.model_dir)?;
        let _ = self.id_alignment.set(alignment);
        Ok(self
            .id_alignment
            .get()
            .expect("id_alignment set in previous statement"))
    }

    /// Encode documents into per-token embeddings AND the token id that
    /// produced each row, aligned 1:1 (`ids[i]` labels `embeddings.row(i)`).
    ///
    /// # Investigation finding
    ///
    /// `next-plaid-onnx::Colbert::encode_documents` does NOT hand back one
    /// row per raw input token id — it rewrites the sequence and filters it:
    ///
    /// - It tokenizes without the doc marker, then inserts the `[D] `
    ///   document-prefix token id at position 1 (right after `[CLS]`),
    ///   matching PyLate's tokenization approach (`next-plaid-onnx` 1.6.2
    ///   `src/lib.rs:1816-1823`, assembled per-document in
    ///   `prepare_batch_from_tokenizer_encodings`, `:2007-2131`).
    /// - When extracting embedding rows from the ONNX output it drops
    ///   padding rows AND rows whose token id is in `onnx_config.json`'s
    ///   `skiplist_words` — a ColBERT-style punctuation filter
    ///   (`encode_prepared_batch_with_session`, `:2154-2246`, gated by the
    ///   `filter_skiplist` flag that `encode_documents`/`encode_documents_raw`
    ///   always pass as `true`, `:1118-1138`).
    ///
    /// So the true output shape is `[CLS], [D], content tokens minus
    /// punctuation, optional [SEP] (only if the document was truncated to
    /// `document_length`)`.
    ///
    /// Nothing in `next-plaid-onnx`'s public API exposes these ids directly:
    /// `Colbert::tokenize_documents` (`:1140-1142`) returns a
    /// `PreparedDocumentBatch` that DOES carry the exact sequence internally
    /// (its private `all_token_ids` field, alongside `original_lengths` and
    /// `filter_skiplist`, `:742-759`), but that struct only exposes
    /// `batch_size()`/`batch_max_len()` publicly (`:766-774`) — there is no
    /// accessor for the ids. So this method replicates the construction
    /// itself from the same `tokenizer.json` + `onnx_config.json` the encoder
    /// loads (see [`DocIdAlignment`] / [`build_doc_token_ids`]) and applies
    /// the identical skiplist filter, rather than trying to read the ids off
    /// the encoder. `ids.len() == embeddings.nrows()` always holds for every
    /// document — callers (e.g. static-token-table distillation) can zip them
    /// directly.
    pub fn encode_documents_with_ids(
        &self,
        texts: &[String],
    ) -> Result<Vec<(Vec<u32>, TokenEmbeddings)>> {
        let embeddings = self.encode_documents(texts)?;
        let alignment = self.id_alignment()?;
        texts
            .iter()
            .zip(embeddings)
            .map(|(text, emb)| {
                let ids = build_doc_token_ids(alignment, text)?;
                anyhow::ensure!(
                    ids.len() == emb.nrows(),
                    "token id/embedding row mismatch ({} ids vs {} rows) for document {:?}",
                    ids.len(),
                    emb.nrows(),
                    text
                );
                Ok((ids, emb))
            })
            .collect()
    }

    /// Vocabulary size of the tokenizer this embedder uses for document
    /// token-id alignment (see [`ColbertEmbedder::encode_documents_with_ids`]).
    /// Used by static-table distillation (Task 3) to size its per-token
    /// accumulators (`vocab_size × dims`) up front.
    ///
    /// Passes `with_added_tokens: true` so the reported size matches the
    /// id-space `encode_documents_with_ids` can actually emit — added
    /// special tokens (like the `[D]` document marker) occupy ids inside
    /// that space, and an accumulator sized without them could be too
    /// small, causing an out-of-bounds write for legitimate ids the encoder
    /// hands back.
    ///
    /// # Errors
    ///
    /// Returns an error only if the tokenizer/config files under
    /// `model_dir` are missing or malformed — the same failure
    /// [`ColbertEmbedder::encode_documents_with_ids`] would hit on its own
    /// first call (`id_alignment()` is shared, memoizing HF-crate
    /// `tokenizer.json` file access with a slow tokenizer, and
    /// `onnx_config.json` parsing). There is no runtime scenario where the
    /// model directory would be valid at construction and then go missing
    /// mid-process; it is not user-editable at runtime, so this is a
    /// construction-time integrity property of the embedder, not something
    /// that "presently is, then later isn't".
    pub fn tokenizer_vocab_size(&self) -> Result<usize> {
        Ok(self.id_alignment()?.tokenizer.get_vocab_size(true))
    }

    /// Number of vocab tokens actually reachable on the document side. With
    /// `do_lower_case: true` (LateOn's onnx_config), any token containing an
    /// uppercase character can never be produced from lowered input —
    /// 25.9% of the LateOn vocab (13,035 of 50,370, empirically measured;
    /// see results/ember-gate1/vocab-coverage-investigation.md, which
    /// independently reports 12,953). Data-driven: when the model config
    /// doesn't lower, this equals `tokenizer_vocab_size`.
    ///
    /// # Implementation note: must check the *decoded* token, not the raw
    /// vocab key
    ///
    /// This tokenizer is byte-level BPE: its raw vocab keys use a
    /// byte-to-unicode remapping (à la GPT-2) where non-printable bytes —
    /// including the extremely common word-boundary space byte — are
    /// stood in for by characters like `Ġ` (U+0120) and `Ċ` (U+010A), which
    /// land in the Latin Extended-A block and are themselves classified as
    /// Unicode-uppercase. Checking `char::is_uppercase` against the raw key
    /// (as opposed to `tokenizer.decode`'s reconstructed text) therefore
    /// misclassifies most word-initial tokens as "contains uppercase" —
    /// empirically 35,225 of 50,370 raw keys vs. only 13,035 of the
    /// correctly decoded tokens for `LateOn-Code-edge` (the latter matches
    /// the investigation doc's 12,953 within expected special-token
    /// variance). Decoding each id is the only correct way to answer "would
    /// lowercasing ever produce this token".
    ///
    /// # This count is an approximation, not an exact reachability oracle
    ///
    /// Byte-level BPE vocabularies also contain **byte-fragment** tokens —
    /// pieces of a multi-byte UTF-8 codepoint that only form valid text when
    /// concatenated with a sibling fragment from the same original
    /// character. Decoded in isolation (as this method does, one id at a
    /// time), such a fragment either renders as a Unicode replacement
    /// character (U+FFFD) — counted as "reachable" with no real basis, since
    /// replacement chars are never case-bearing — or `decode` returns an
    /// error, which `.is_ok_and(...)` silently treats as "not reachable"
    /// even though the fragment may in fact be produced by real tokenization
    /// (just never standalone). Neither outcome is validated against actual
    /// document-side tokenization. The returned count is therefore a
    /// **diagnostic approximation** — accurate enough to distinguish "most
    /// of the vocab is reachable" from "half of it is" (the bug this method
    /// replaced), but not a byte-for-byte guarantee for every id. Do not
    /// treat it as an exact reachability oracle.
    ///
    /// # Errors
    ///
    /// Same failure mode as [`ColbertEmbedder::tokenizer_vocab_size`]: only if
    /// `tokenizer.json`/`onnx_config.json` under `model_dir` are missing or
    /// malformed.
    pub fn tokenizer_reachable_vocab_count(&self) -> Result<usize> {
        let alignment = self.id_alignment()?;
        if !alignment.do_lower_case {
            return Ok(alignment.tokenizer.get_vocab_size(true));
        }
        Ok(alignment
            .tokenizer
            .get_vocab(true)
            .values()
            .filter(|&&id| {
                alignment
                    .tokenizer
                    .decode(&[id], false)
                    .is_ok_and(|decoded| !decoded.chars().any(char::is_uppercase))
            })
            .count())
    }

    /// Returns true if the ONNX session has been materialized.
    /// Test/diagnostics helper — not used in production paths.
    #[doc(hidden)]
    pub fn is_initialized(&self) -> bool {
        self.encoder.get().is_some()
    }

    /// Returns how many times `build_encoder` has been invoked. Test helper for
    /// the lazy-init race-condition test (Finding 10) — should always be ≤ 1
    /// even under concurrent first-callers.
    #[doc(hidden)]
    pub fn build_count(&self) -> usize {
        self.build_count.load(std::sync::atomic::Ordering::Relaxed)
    }
}

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

    /// Locate the local LateOn-Code-edge ColBERT model directory used by the
    /// model-gated encode test below. Returns `None` (causing the test to
    /// skip) when the model hasn't been downloaded, so `cargo test` stays
    /// green in environments without it (CI, fresh clones). Never downloads
    /// on its own — mirrors the path `model_manager::ensure_colbert_model`
    /// downloads to, resolved the same repo-agnostic way production code
    /// does (`SemantexConfig::models_dir`), never a hardcoded path.
    fn test_model_dir() -> Option<PathBuf> {
        let dir = crate::config::SemantexConfig::default()
            .models_dir()
            .join("LateOn-Code-edge");
        dir.join("model_int8.onnx").exists().then_some(dir)
    }

    #[test]
    #[ignore = "loads the ColBERT model and builds a real ONNX session; run with --ignored"]
    fn encode_with_ids_rows_align() {
        let Some(model_dir) = test_model_dir() else {
            return;
        }; // follow existing gate helper
        let e = ColbertEmbedder::new(&model_dir).unwrap();
        let texts = vec!["fn main() { println!(\"hi\"); }".to_string()];
        let out = e.encode_documents_with_ids(&texts).unwrap();
        assert_eq!(out.len(), 1);
        let (ids, emb) = &out[0];
        assert_eq!(ids.len(), emb.nrows(), "one token id per embedding row");
        assert!(!ids.is_empty());
        // Consistency with the plain path: same matrix.
        let plain = e.encode_documents(&texts).unwrap();
        assert_eq!(emb, &plain[0]);
    }

    /// Covers the truncation branch in `assemble_and_filter_ids`
    /// (`content_prefix_len` / `keep_sep`), which the short document in
    /// `encode_with_ids_rows_align` never exercises.
    ///
    /// This is a PURE unit test on synthetic ids, not an end-to-end encode —
    /// deliberately. `LateOn-Code-edge`'s `tokenizer.json` embeds its own
    /// truncation at exactly `document_length - 1` (2047), which the
    /// `tokenizers` crate applies inside `.encode()` before our code ever
    /// sees the ids — so no real input text can make
    /// `alignment.tokenizer.encode(...)` hand back more than 2047 raw ids for
    /// this model, and the `real_len > truncate_limit` branch is unreachable
    /// via the public `encode_documents_with_ids` API today (verified: a
    /// 3000-distinct-identifier document was tried first and came back
    /// pre-truncated to exactly 2047 raw ids by the tokenizer itself). Rather
    /// than leave the branch untested because of that, `assemble_and_filter_ids`
    /// was factored out of `build_doc_token_ids` specifically so this branch
    /// can be exercised directly with a synthetic ids sequence, independent
    /// of any tokenizer's own truncation policy.
    #[test]
    fn truncated_ids_keep_trailing_sep_and_stay_aligned() {
        const CLS: u32 = 1;
        const SEP: u32 = 2;
        const DOC_MARKER: u32 = 999;
        const DOCUMENT_LENGTH: usize = 10; // truncate_limit = 9

        // [CLS], 5000 distinct content ids, [SEP] — far longer than
        // DOCUMENT_LENGTH, forcing the `real_len > truncate_limit` branch.
        let content: Vec<u32> = (10..5010).collect();
        let mut ids = vec![CLS];
        ids.extend(&content);
        ids.push(SEP);

        let skiplist_ids = HashSet::new();
        let out = assemble_and_filter_ids(DOC_MARKER, DOCUMENT_LENGTH, &skiplist_ids, &ids);

        // keep_sep path: content_prefix_len = truncate_limit - 1 = 8, so the
        // final sequence is [CLS] + [D] + 7 content ids + [SEP] = 10 =
        // DOCUMENT_LENGTH exactly.
        assert_eq!(
            out.len(),
            DOCUMENT_LENGTH,
            "truncated sequence must land exactly at document_length"
        );
        assert_eq!(out[0], CLS, "[CLS] must stay first");
        assert_eq!(out[1], DOC_MARKER, "[D] marker inserted at position 1");
        assert_eq!(
            &out[2..out.len() - 1],
            &content[..DOCUMENT_LENGTH - 3],
            "kept content must be a PREFIX of the original content (truncation drops the tail)"
        );
        assert_eq!(
            *out.last().unwrap(),
            SEP,
            "truncated document must keep the trailing [SEP], not cut it off mid-content"
        );
    }

    /// LateOn's `onnx_config.json` sets `do_lower_case: true`, so any vocab
    /// token containing an uppercase character can never be produced from
    /// lowered document input — `tokenizer_reachable_vocab_count` must report
    /// strictly fewer tokens than the raw vocab for this model (see
    /// `results/ember-gate1/vocab-coverage-investigation.md`).
    #[test]
    #[ignore = "loads the ColBERT model and builds a real ONNX session; run with --ignored"]
    fn reachable_vocab_is_smaller_than_raw_vocab_for_lowercasing_model() {
        let Some(model_dir) = test_model_dir() else {
            return;
        };
        let e = ColbertEmbedder::new(&model_dir).unwrap();
        let raw = e.tokenizer_vocab_size().unwrap();
        let reachable = e.tokenizer_reachable_vocab_count().unwrap();
        assert!(
            reachable < raw,
            "LateOn lowers input: reachable {reachable} must be < raw {raw}"
        );
        assert!(
            reachable > raw / 2,
            "sanity: ~74% of vocab is lowercase-only (reachable={reachable}, raw={raw})"
        );
    }

    /// End-to-end companion to the pure unit test above: confirms a real,
    /// genuinely long document still produces aligned ids/embeddings through
    /// the full `encode_documents_with_ids` pipeline — the scenario Task 3's
    /// distillation corpus will actually hit. It does NOT exercise the
    /// `keep_sep` branch (see the note on `build_doc_token_ids`: this
    /// model's tokenizer pre-truncates to exactly `document_length - 1`
    /// itself), but it does prove alignment holds once the tokenizer's own
    /// truncation kicks in, and that [SEP] still ends up last via the
    /// ordinary (non-`keep_sep`) path.
    #[test]
    #[ignore = "loads the ColBERT model and builds a real ONNX session; run with --ignored"]
    fn encode_with_ids_stays_aligned_for_a_long_document() {
        let Some(model_dir) = test_model_dir() else {
            return;
        };
        let e = ColbertEmbedder::new(&model_dir).unwrap();

        // Each `identifier_N` is distinct so it can't collapse into a single
        // existing vocab token; 3000 of them is far more than this model's
        // tokenizer needs to hit its own embedded truncation.
        let long_text: String = (0..3000)
            .map(|i| format!("identifier_{i}"))
            .collect::<Vec<_>>()
            .join(" ");
        let texts = vec![long_text];

        let out = e.encode_documents_with_ids(&texts).unwrap();
        assert_eq!(out.len(), 1);
        let (ids, emb) = &out[0];
        assert_eq!(
            ids.len(),
            emb.nrows(),
            "one token id per embedding row, even for a long/truncated document"
        );
        assert!(!ids.is_empty());

        let alignment = e.id_alignment().unwrap();

        // Confirm the fixture actually reached this model's truncation limit
        // (tokenizer.json's own embedded max_length), not that it happened
        // to tokenize short.
        let raw = alignment.tokenizer.encode(texts[0].as_str(), true).unwrap();
        assert_eq!(
            raw.get_ids().len(),
            alignment.document_length - 1,
            "tokenizer.json's embedded truncation should cap raw ids at document_length - 1"
        );

        let sep_id = alignment
            .tokenizer
            .token_to_id("[SEP]")
            .expect("[SEP] must be in the tokenizer vocabulary");
        assert_eq!(
            *ids.last().unwrap(),
            sep_id,
            "long document must still end with [SEP]"
        );
    }

    #[test]
    fn default_index_threads_uses_all_cores_when_only_one_build_slot_exists() {
        // A box that can only ever run one full build at a time (RAM-limited)
        // has no sibling build to leave cores for — use them all (clamped).
        assert_eq!(default_index_threads(4, 1), 4);
        assert_eq!(default_index_threads(2, 1), 2, "clamped to floor of 2");
        assert_eq!(default_index_threads(16, 1), 8, "clamped to ceiling of 8");
        assert_eq!(default_index_threads(1, 1), 2, "clamped to floor of 2");
    }

    #[test]
    fn default_index_threads_halves_cores_when_multiple_build_slots_exist() {
        // Bigger boxes where several builds could run at once keep the
        // conservative half-the-cores split (unchanged legacy behaviour).
        assert_eq!(default_index_threads(8, 2), 4);
        assert_eq!(default_index_threads(32, 4), 8, "clamped to ceiling of 8");
        assert_eq!(default_index_threads(4, 2), 2, "clamped to floor of 2");
    }

    /// `ColbertEmbedder::new()` should not error or build the ONNX session when
    /// the model directory exists but the model files are missing/corrupt —
    /// the lazy init contract is "fail late, never at construction".
    #[test]
    fn new_is_lazy_does_not_build_session() {
        let tmp = tempfile::TempDir::new().unwrap();
        // Directory exists, but no model files — should still succeed at construction.
        let embedder = ColbertEmbedder::new(tmp.path())
            .expect("constructor should succeed for any existing directory");
        assert!(
            !embedder.is_initialized(),
            "ONNX session must not be materialized at construction time"
        );
    }

    /// `ColbertEmbedder::new()` must reject non-existent model dirs early,
    /// since the alternative (deferring the error to first encode) would
    /// hide configuration bugs until query time.
    #[test]
    fn new_rejects_missing_model_dir() {
        let res = ColbertEmbedder::new(Path::new("/nonexistent/path/that/does/not/exist"));
        assert!(res.is_err(), "missing model dir must fail at construction");
    }

    /// Finding 10: when multiple threads simultaneously call `encoder()` for the
    /// first time, only ONE may invoke `build_encoder`. The naive
    /// check-then-build pattern previously let N threads each build a ~150 MB
    /// ONNX session and silently drop N-1 of them, doubling peak RSS at the
    /// worst possible moment (E8c warm-up colliding with the first user query).
    ///
    /// We cannot load a real ONNX model in unit tests, so this test exercises
    /// the lock contract end-to-end with a simulated "successful build" via a
    /// shared barrier and explicit OnceLock seeding from inside the build_lock
    /// region. The invariant is:
    ///
    ///   "After one thread has set the OnceLock under the build_lock, no
    ///    subsequent waiter that observes the lock-release will call
    ///    `build_encoder` again."
    ///
    /// We verify this by checking `build_count()` after all racing callers
    /// complete: at most one invocation total, regardless of thread count.
    ///
    /// Without the fix this test consistently observes 2-4 build invocations
    /// (each thread that passed the initial `encoder.get()` check still races
    /// to call `build_encoder`). With the fix it observes exactly 1.
    #[test]
    fn encoder_init_is_serialized_under_concurrency() {
        use std::sync::Arc;
        use std::sync::Barrier;
        use std::sync::atomic::{AtomicUsize, Ordering};

        // We test the lock semantics by inlining the same check-build-set
        // pattern the production code uses, but with an instrumented "build"
        // that simply counts invocations. This proves the locking discipline
        // without needing a real ~150MB ONNX session.
        struct TestEmbedder {
            encoder: OnceLock<u32>,
            build_lock: std::sync::Mutex<()>,
            build_count: AtomicUsize,
        }

        impl TestEmbedder {
            fn new() -> Self {
                Self {
                    encoder: OnceLock::new(),
                    build_lock: std::sync::Mutex::new(()),
                    build_count: AtomicUsize::new(0),
                }
            }

            /// Mirrors the production `encoder()` discipline exactly.
            fn get_or_init(&self) -> &u32 {
                if let Some(v) = self.encoder.get() {
                    return v;
                }
                let _guard = self
                    .build_lock
                    .lock()
                    .unwrap_or_else(std::sync::PoisonError::into_inner);
                if let Some(v) = self.encoder.get() {
                    return v;
                }
                // Simulated build — counted, but cheap.
                self.build_count.fetch_add(1, Ordering::Relaxed);
                // Sleep a bit so other threads have time to pile up at the
                // lock; this widens the race window dramatically and would
                // expose the bug immediately in the unfixed version.
                std::thread::sleep(std::time::Duration::from_millis(50));
                let _ = self.encoder.set(42);
                self.encoder.get().expect("encoder set above")
            }
        }

        let embedder = Arc::new(TestEmbedder::new());
        let n_threads = 8;
        let barrier = Arc::new(Barrier::new(n_threads));
        let mut handles = Vec::with_capacity(n_threads);
        for _ in 0..n_threads {
            let emb = Arc::clone(&embedder);
            let b = Arc::clone(&barrier);
            handles.push(std::thread::spawn(move || {
                b.wait();
                let v = emb.get_or_init();
                assert_eq!(*v, 42);
            }));
        }
        for h in handles {
            h.join().unwrap();
        }

        let count = embedder.build_count.load(Ordering::Relaxed);
        assert_eq!(
            count, 1,
            "build path must be invoked exactly once under {n_threads} concurrent \
             first-callers (observed {count}). If this fails, the check-build-set \
             pattern in `ColbertEmbedder::encoder` has regressed — see Finding 10."
        );
    }

    /// Confirms the production `ColbertEmbedder` has the `build_lock` field and
    /// preserves the read-fast-path. After a successful `set`, subsequent
    /// `encoder()` calls must NOT increment `build_count`.
    #[test]
    fn build_count_does_not_grow_on_cached_reads() {
        // We can't trigger a real successful build without a model, so we
        // assert the structural property: the counter starts at zero, and the
        // accessor exists. The end-to-end fast-path is exercised by the
        // production e2e smoke test (ColBERT race test in the task spec).
        let tmp = tempfile::TempDir::new().unwrap();
        let embedder = ColbertEmbedder::new(tmp.path()).unwrap();
        assert_eq!(embedder.build_count(), 0);
        assert!(!embedder.is_initialized());
    }
}