Skip to main content

nabled_embeddings/
lib.rs

1//! Lightweight, ndarray-native, Arrow-zero-copy compute and rerank layer for embedding vectors.
2//!
3//! Bring vectors from any model, compute exactly, deploy anywhere. `nabled-embeddings` composes
4//! existing `nabled-linalg` vector kernels and `nabled-ml` PCA into the post-embedding numerics a
5//! retrieval pipeline needs: normalize, score, rerank, exact brute-force kNN, and PCA compression.
6//!
7//! # What this is (and is not)
8//!
9//! This is the **exact rerank / compute layer that sits next to a vector store**, not a vector
10//! database. It is deliberately:
11//!
12//! - **bring-your-own-vectors:** it accepts `Array2<T>` / `ArrayView2<T>` (`f32`/`f64` via
13//!   [`NabledReal`](nabled_core::scalar::NabledReal)). Any encoder's dense float vectors plug in
14//!   unchanged because the crate depends only on shape and dtype, never the model.
15//! - **numerics-only:** no model inference, no tokenizers, no ONNX/candle.
16//! - **Arrow-free in the core:** Arrow/Lance ingress lives at the `nabled::arrow` facade and in
17//!   `pynabled`, never here.
18//!
19//! Non-goals: it is **not** an ANN/vector-search engine (no HNSW/IVF index), **not** storage, and
20//! **not** a FAISS/HNSW replacement. Use an ANN index for recall, then rerank its candidates here.
21//!
22//! # Two rules the math cannot enforce
23//!
24//! 1. Query and corpus must come from the **same model** and share the same `dim`.
25//! 2. You pick the [`Metric`] to match how the model was trained: [`Metric::Cosine`] (the default),
26//!    [`Metric::Dot`] for maximum-inner-product (MIPS) models, or [`Metric::L2`] where applicable.
27//!    Dot on un-normalized vectors favors larger-norm rows by design.
28//!
29//! # Feature flags
30//!
31//! 1. `blas`: forwards BLAS acceleration to `nabled-linalg`/`nabled-ml`.
32//! 2. `lapack-provider`: enables provider-dispatched PCA paths.
33//! 3. `openblas-system`/`openblas-static`/`netlib-system`/`netlib-static`/`magma-system`: provider
34//!    backends forwarded to the lower crates.
35//!
36//! # Quick example
37//!
38//! ```rust
39//! use nabled_embeddings::{Metric, rerank};
40//! use ndarray::arr2;
41//!
42//! let corpus = arr2(&[[1.0_f64, 0.0], [0.0, 1.0], [0.9, 0.1]]);
43//! let query = arr2(&[[1.0_f64, 0.0]]);
44//! let top = rerank(&query.row(0), &corpus.view(), 2, Metric::Cosine)?;
45//! assert_eq!(top[0].index, 0);
46//! # Ok::<(), nabled_embeddings::EmbeddingError>(())
47//! ```
48
49pub mod cache;
50pub mod compress;
51pub mod error;
52pub mod knn;
53pub mod metrics;
54pub mod mmr;
55pub mod normalize;
56pub mod quantize;
57pub mod rerank;
58pub mod similarity;
59pub mod topk;
60
61pub use cache::CorpusWorkspace;
62pub use compress::{PcaModel, compress, compress_into, compress_view, fit_pca};
63pub use error::EmbeddingError;
64pub use knn::brute_force_knn;
65pub use metrics::{mean_reciprocal_rank, ndcg_at_k, recall_at_k, reciprocal_rank};
66pub use mmr::mmr;
67pub use normalize::{normalize_rows, normalize_rows_into, normalize_rows_view};
68pub use quantize::{QuantizedMatrix, dequantize, quantize_rows};
69pub use rerank::{batch_rerank_with_ids, rerank, rerank_with_ids};
70pub use similarity::{
71    Metric, query_corpus_scores, query_corpus_scores_into, query_corpus_scores_view,
72};
73pub use topk::{Neighbor, NeighborWithId, attach_ids, top_k};