Skip to main content

next_plaid/
lib.rs

1//! Next-Plaid: CPU-based PLAID implementation for multi-vector search
2//!
3//! This crate provides a pure Rust, CPU-only implementation of the PLAID algorithm
4//! for efficient multi-vector search (late interaction retrieval).
5
6// Link BLAS implementation when feature is enabled
7#[cfg(feature = "accelerate")]
8extern crate blas_src;
9
10#[cfg(feature = "openblas")]
11extern crate openblas_src;
12
13pub mod codec;
14#[cfg(feature = "cuda")]
15pub mod cuda;
16pub mod delete;
17pub mod embeddings;
18pub mod error;
19pub mod filtering;
20pub mod index;
21pub mod kmeans;
22pub mod maxsim;
23pub mod mmap;
24pub mod search;
25pub mod text_search;
26pub mod update;
27pub mod utils;
28
29pub use codec::ResidualCodec;
30pub use delete::delete_from_index;
31pub use error::{Error, Result};
32pub use index::MmapIndex;
33pub use index::{
34    encode_index_chunk, prepare_codec_artifacts, write_index_from_encoded_chunks,
35    EncodedIndexChunk, IndexConfig, Metadata, PreparedCodecArtifacts,
36};
37pub use kmeans::{
38    compute_centroids, compute_centroids_from_documents, compute_kmeans, estimate_num_partitions,
39    ComputeKmeansConfig, FastKMeans, KMeansConfig,
40};
41pub use search::{QueryResult, SearchParameters};
42pub use text_search::FtsTokenizer;
43pub use update::UpdateConfig;
44
45#[cfg(feature = "cuda")]
46pub use cuda::{clear_cuda_broken, is_cuda_broken, mark_cuda_broken, CudaContext};
47
48/// Check if GPU-only mode is forced via environment variable.
49/// Only checks the canonical `NEXT_PLAID_FORCE_GPU` env var.
50/// The higher-level `colgrep` crate's `apply_acceleration_mode()` propagates
51/// CLI flags and `COLGREP_*`/`FORCE_*` vars into this canonical var.
52pub fn is_force_gpu() -> bool {
53    std::env::var("NEXT_PLAID_FORCE_GPU")
54        .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
55        .unwrap_or(false)
56}
57
58/// Check if CPU-only mode is forced via environment variable.
59/// Only checks the canonical `NEXT_PLAID_FORCE_CPU` env var.
60pub fn is_force_cpu() -> bool {
61    !is_force_gpu()
62        && std::env::var("NEXT_PLAID_FORCE_CPU")
63            .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
64            .unwrap_or(false)
65}