lc_embeddings/local/
mod.rs

1// lc-embeddings/src/local/mod.rs
2//! Local embedding implementations
3//!
4//! Contains two implementations:
5//! - `BagOfWordsEmbeddings`: Lightweight word-frequency hash embedding (pure Rust, no external deps), always available
6//! - `LocalEmbeddings`: ONNX Runtime-based neural network embedding (requires `local-embeddings` feature)
7//!
8//! `BagOfWordsEmbeddings` is suitable for offline, privacy, zero-cost coarse-grained retrieval;
9//! `LocalEmbeddings` is suitable for high-quality semantic embedding scenarios (e.g., BGE/E5 models).
10
11use async_trait::async_trait;
12
13#[cfg(feature = "local-embeddings")]
14use std::path::Path;
15
16use crate::{EmbeddingError, Embeddings};
17
18// ---------------------------------------------------------------------------
19// BagOfWordsEmbeddings — word-frequency hash + L2 normalization (always available)
20// ---------------------------------------------------------------------------
21
22/// Lightweight local embedding (word-frequency hash + L2 normalization)
23///
24/// Based on word frequency + hashing, no API calls, suitable for offline, privacy, zero-cost coarse-grained retrieval.
25///
26/// Note: This is a lightweight implementation (bag-of-words hash) with limited semantic quality;
27/// for high-quality neural network embeddings (BGE/E5 via `ort`), enable the `local-embeddings`
28/// feature and use the `LocalEmbeddings` type (only exists under that feature).
29pub struct BagOfWordsEmbeddings {
30    dim: usize,
31}
32
33impl BagOfWordsEmbeddings {
34    /// Create local embedding with specified dimension
35    pub fn new(dim: usize) -> Self {
36        Self { dim: dim.max(1) }
37    }
38
39    /// Default dimension 256
40    pub fn default_dim() -> Self {
41        Self::new(256)
42    }
43
44    /// Tokenize: English by non-alphanumeric split (lowercased), Chinese/non-ASCII by single character
45    fn tokenize(text: &str) -> Vec<String> {
46        let mut tokens = Vec::new();
47        let mut current = String::new();
48        for c in text.chars() {
49            if c.is_alphanumeric() {
50                if c.is_ascii() {
51                    current.push(c.to_ascii_lowercase());
52                } else {
53                    // Non-ASCII (Chinese etc.) single character as token
54                    if !current.is_empty() {
55                        tokens.push(std::mem::take(&mut current));
56                    }
57                    tokens.push(c.to_string());
58                }
59            } else if !current.is_empty() {
60                tokens.push(std::mem::take(&mut current));
61            }
62        }
63        if !current.is_empty() {
64            tokens.push(current);
65        }
66        tokens
67    }
68
69    /// FNV-1a hash
70    fn hash(s: &str) -> u64 {
71        let mut h: u64 = 0xcbf29ce484222325;
72        for b in s.bytes() {
73            h ^= b as u64;
74            h = h.wrapping_mul(0x100000001b3);
75        }
76        h
77    }
78
79    /// Compute embedding vector (word-frequency hash + L2 normalization)
80    fn embed(&self, text: &str) -> Vec<f32> {
81        let mut v = vec![0.0f32; self.dim];
82        for token in Self::tokenize(text) {
83            let idx = (Self::hash(&token) as usize) % self.dim;
84            v[idx] += 1.0;
85        }
86        // L2 normalization
87        let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
88        if norm > 0.0 {
89            for x in &mut v {
90                *x /= norm;
91            }
92        }
93        v
94    }
95}
96
97impl Default for BagOfWordsEmbeddings {
98    fn default() -> Self {
99        Self::default_dim()
100    }
101}
102
103#[async_trait]
104impl Embeddings for BagOfWordsEmbeddings {
105    async fn embed_query(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
106        if text.trim().is_empty() {
107            return Err(EmbeddingError::EmptyInput);
108        }
109        Ok(self.embed(text))
110    }
111
112    fn dimension(&self) -> usize {
113        self.dim
114    }
115
116    fn model_name(&self) -> &str {
117        "local-bow"
118    }
119}
120
121// ---------------------------------------------------------------------------
122// LocalEmbeddings — ONNX Runtime neural network embedding (requires local-embeddings feature)
123// ---------------------------------------------------------------------------
124
125#[cfg(feature = "local-embeddings")]
126mod nn;
127
128// When local-embeddings feature is enabled, re-export LocalEmbeddings and its builder
129#[cfg(feature = "local-embeddings")]
130pub use nn::{LocalEmbeddings, LocalEmbeddingsBuilder};
131
132// ---------------------------------------------------------------------------
133// 1.0:LocalEmbeddings 降级别名已移除
134// ---------------------------------------------------------------------------
135
136// 1.0 起,无 `local-embeddings` feature 时不再提供 `LocalEmbeddings` 名字(原
137// `BagOfWordsEmbeddings` 降级别名)。使用者被迫显式选边:要么 `BagOfWordsEmbeddings`
138// (词袋哈希),要么开启 feature 用 ONNX 神经嵌入——彻底封掉"以为在用语义向量、
139// 实际是词频"的坑。有 feature 时 `LocalEmbeddings` 为 nn 模块的 ONNX 实现
140// (见上方 `#[cfg(feature = "local-embeddings")] pub use nn::...`)。
141
142// ---------------------------------------------------------------------------
143// Tests
144// ---------------------------------------------------------------------------
145
146#[cfg(test)]
147mod tests;
lc_embeddings/local/mod.rs

lc_embeddings/local/
mod.rs
