lattice_embed/lib.rs
1//! **Stability tier**: Unstable
2//!
3//! The SIMD dispatch layer (`simd/`) and model-loading API are actively evolving.
4//! Platform consumers should use this crate via `lattice-engine`, not directly.
5//! The 21 `unsafe` blocks are gated SIMD intrinsic calls (AVX512/AVX2/NEON); the
6//! 1 `dead_code_allows` retains a superseded dot-product fallback for reference.
7//! See `foundation/STABILITY.md` for the full policy.
8//!
9//! # lattice-embed
10
11#![warn(missing_docs)]
12#![allow(clippy::clone_on_copy)]
13//!
14//! Vector embedding generation with SIMD-accelerated operations for the lattice-runtime substrate.
15//!
16//! This crate provides embedding generation services that convert text into
17//! high-dimensional vector representations suitable for semantic search and
18//! similarity matching.
19//!
20//! ## Features
21//!
22//! - **Native Embeddings**: Generate embeddings locally using pure Rust inference (default)
23//! - **BGE Models**: Support for BGE family of models (small/base/large)
24//! - **Async API**: Full async/await support with tokio
25//! - **SIMD Acceleration**: AVX2/NEON optimized vector operations
26//!
27//! ## Quick Start
28//!
29//! ```rust,no_run
30//! use lattice_embed::{EmbeddingService, EmbeddingModel, NativeEmbeddingService};
31//!
32//! #[tokio::main]
33//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
34//! let service = NativeEmbeddingService::default();
35//!
36//! let embedding = service.embed_one(
37//! "The quick brown fox jumps over the lazy dog",
38//! EmbeddingModel::default(),
39//! ).await?;
40//!
41//! println!("Embedding dimension: {}", embedding.len());
42//! // Output: Embedding dimension: 384
43//!
44//! Ok(())
45//! }
46//! ```
47//!
48//! ## Batch Processing
49//!
50//! ```rust,no_run
51//! use lattice_embed::{EmbeddingService, EmbeddingModel, NativeEmbeddingService};
52//!
53//! #[tokio::main]
54//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
55//! let service = NativeEmbeddingService::default();
56//!
57//! let texts = vec![
58//! "First document".to_string(),
59//! "Second document".to_string(),
60//! "Third document".to_string(),
61//! ];
62//!
63//! let embeddings = service.embed(&texts, EmbeddingModel::BgeSmallEnV15).await?;
64//! assert_eq!(embeddings.len(), 3);
65//!
66//! Ok(())
67//! }
68//! ```
69//!
70//! ## Available Models
71//!
72//! | Model | Dimensions | Use Case |
73//! |-------|------------|----------|
74//! | `BgeSmallEnV15` | 384 | Fast, general purpose (default) |
75//! | `BgeBaseEnV15` | 768 | Balanced quality/speed |
76//! | `BgeLargeEnV15` | 1024 | Highest quality |
77
78pub mod backfill;
79mod cache;
80mod error;
81pub mod migration;
82mod model;
83mod service;
84pub mod simd;
85pub mod types;
86
87pub use cache::{CacheStats, DEFAULT_CACHE_CAPACITY, EmbeddingCache, ShardStats};
88pub use error::{EmbedError, Result};
89pub use model::{EmbeddingModel, MIN_MRL_OUTPUT_DIM, ModelConfig, ModelProvenance};
90pub use service::{DEFAULT_MAX_BATCH_SIZE, EmbeddingService, MAX_TEXT_CHARS};
91pub use simd::{SimdConfig, simd_config};
92
93#[cfg(feature = "native")]
94pub use service::{CachedEmbeddingService, NativeEmbeddingService};
95
96/// Utility functions for vector operations.
97///
98/// All functions in this module are SIMD-accelerated when available (AVX2 on x86_64, NEON on aarch64).
99/// Runtime feature detection ensures automatic fallback to scalar implementations
100/// on systems without SIMD support.
101pub mod utils {
102 use crate::simd;
103
104 /// **Stable**: external consumers may depend on this; breaking changes require a SemVer bump.
105 ///
106 /// Compute cosine similarity between two vectors.
107 ///
108 /// Uses SIMD acceleration (AVX2/NEON) when available, with automatic scalar fallback.
109 ///
110 /// Returns a value between -1.0 and 1.0, where 1.0 indicates
111 /// identical direction and -1.0 indicates opposite direction.
112 ///
113 /// # Performance
114 ///
115 /// | Dimension | Scalar | SIMD |
116 /// |-----------|--------|------|
117 /// | 384 | ~650ns | ~90ns |
118 /// | 768 | ~1300ns | ~180ns |
119 /// | 1024 | ~1700ns | ~240ns |
120 ///
121 /// # Example
122 ///
123 /// ```rust
124 /// use lattice_embed::utils::cosine_similarity;
125 ///
126 /// let a = vec![1.0, 0.0, 0.0];
127 /// let b = vec![1.0, 0.0, 0.0];
128 /// assert!((cosine_similarity(&a, &b) - 1.0).abs() < 0.0001);
129 ///
130 /// let c = vec![1.0, 0.0, 0.0];
131 /// let d = vec![0.0, 1.0, 0.0];
132 /// assert!((cosine_similarity(&c, &d) - 0.0).abs() < 0.0001);
133 /// ```
134 #[inline]
135 pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
136 simd::cosine_similarity(a, b)
137 }
138
139 /// **Stable**: external consumers may depend on this; breaking changes require a SemVer bump.
140 ///
141 /// Compute dot product between two vectors.
142 ///
143 /// For normalized vectors (like embeddings), this equals cosine similarity.
144 /// Using `dot_product` directly on pre-normalized vectors is faster than
145 /// `cosine_similarity` since it skips norm computation.
146 ///
147 /// # Example
148 ///
149 /// ```rust
150 /// use lattice_embed::utils::dot_product;
151 ///
152 /// let a = vec![1.0, 2.0, 3.0];
153 /// let b = vec![4.0, 5.0, 6.0];
154 /// assert!((dot_product(&a, &b) - 32.0).abs() < 0.0001);
155 /// ```
156 #[inline]
157 pub fn dot_product(a: &[f32], b: &[f32]) -> f32 {
158 simd::dot_product(a, b)
159 }
160
161 /// **Stable**: external consumers may depend on this; breaking changes require a SemVer bump.
162 ///
163 /// Normalize a vector to unit length (L2 normalization).
164 ///
165 /// Modifies the vector in place. After normalization, the vector
166 /// will have magnitude 1.0.
167 ///
168 /// # Example
169 ///
170 /// ```rust
171 /// use lattice_embed::utils::normalize;
172 ///
173 /// let mut v = vec![3.0, 4.0];
174 /// normalize(&mut v);
175 ///
176 /// let magnitude: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
177 /// assert!((magnitude - 1.0).abs() < 0.0001);
178 /// ```
179 #[inline]
180 pub fn normalize(vector: &mut [f32]) {
181 simd::normalize(vector)
182 }
183
184 /// **Stable**: external consumers may depend on this; breaking changes require a SemVer bump.
185 ///
186 /// Compute Euclidean distance between two vectors.
187 ///
188 /// # Example
189 ///
190 /// ```rust
191 /// use lattice_embed::utils::euclidean_distance;
192 ///
193 /// let a = vec![0.0, 0.0];
194 /// let b = vec![3.0, 4.0];
195 /// assert!((euclidean_distance(&a, &b) - 5.0).abs() < 0.0001);
196 /// ```
197 #[inline]
198 pub fn euclidean_distance(a: &[f32], b: &[f32]) -> f32 {
199 simd::euclidean_distance(a, b)
200 }
201
202 /// **Stable**: external consumers may depend on this; breaking changes require a SemVer bump.
203 ///
204 /// Compute cosine similarities for many vector pairs (batch operation).
205 ///
206 /// More efficient than calling `cosine_similarity` in a loop.
207 #[inline]
208 pub fn batch_cosine_similarity(pairs: &[(&[f32], &[f32])]) -> Vec<f32> {
209 simd::batch_cosine_similarity(pairs)
210 }
211
212 /// **Stable**: external consumers may depend on this; breaking changes require a SemVer bump.
213 ///
214 /// Compute dot products for many vector pairs (batch operation).
215 #[inline]
216 pub fn batch_dot_product(pairs: &[(&[f32], &[f32])]) -> Vec<f32> {
217 simd::batch_dot_product(pairs)
218 }
219}
220
221#[cfg(test)]
222mod tests {
223 use super::*;
224
225 #[test]
226 fn test_cosine_similarity_identical() {
227 let a = vec![1.0, 2.0, 3.0];
228 let b = vec![1.0, 2.0, 3.0];
229 let sim = utils::cosine_similarity(&a, &b);
230 assert!((sim - 1.0).abs() < 0.0001);
231 }
232
233 #[test]
234 fn test_cosine_similarity_orthogonal() {
235 let a = vec![1.0, 0.0];
236 let b = vec![0.0, 1.0];
237 let sim = utils::cosine_similarity(&a, &b);
238 assert!(sim.abs() < 0.0001);
239 }
240
241 #[test]
242 fn test_cosine_similarity_opposite() {
243 let a = vec![1.0, 0.0];
244 let b = vec![-1.0, 0.0];
245 let sim = utils::cosine_similarity(&a, &b);
246 assert!((sim + 1.0).abs() < 0.0001);
247 }
248
249 #[test]
250 fn test_normalize() {
251 let mut v = vec![3.0, 4.0];
252 utils::normalize(&mut v);
253 let magnitude: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
254 assert!((magnitude - 1.0).abs() < 0.0001);
255 }
256
257 #[test]
258 fn test_euclidean_distance() {
259 let a = vec![0.0, 0.0, 0.0];
260 let b = vec![1.0, 0.0, 0.0];
261 let dist = utils::euclidean_distance(&a, &b);
262 assert!((dist - 1.0).abs() < 0.0001);
263 }
264
265 #[test]
266 fn test_model_default() {
267 let model = EmbeddingModel::default();
268 assert_eq!(model.dimensions(), 384);
269 }
270}