mindcore 0.2.0

Pluggable, feature-gated memory engine for AI agent applications
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232

//! CandleNativeBackend: granite-small-r2 via ModernBERT
//!
//! Only compiled when `local-embeddings` feature is enabled.
//! Provides 384-dimensional embeddings with 8K token context.

#[cfg(feature = "local-embeddings")]
mod inner {
    use candle_core::{DType, Device, Tensor};
    use candle_nn::VarBuilder;
    use candle_transformers::models::modernbert::{Config, ModernBert};
    use hf_hub::api::sync::Api;
    use hf_hub::{Repo, RepoType};
    use std::path::Path;
    use tokenizers::{PaddingParams, PaddingStrategy, Tokenizer};

    use crate::embeddings::EmbeddingBackend;
    use crate::error::{MindCoreError, Result};

    const MODEL_REPO: &str = "ibm-granite/granite-embedding-small-english-r2";
    const MODEL_NAME: &str = "granite-embedding-small-english-r2";
    const DIMENSIONS: usize = 384;

    /// Embedding backend using IBM granite-small-r2 via candle's ModernBERT.
    ///
    /// Auto-downloads model files (~95MB) from HuggingFace on first use.
    /// Cached at the HuggingFace cache directory (`~/.cache/huggingface/hub/`).
    pub struct CandleNativeBackend {
        model: ModernBert,
        tokenizer: Tokenizer,
        device: Device,
        dimensions_override: Option<usize>,
    }

    impl CandleNativeBackend {
        /// Create with auto-downloaded model from HuggingFace.
        pub fn new() -> Result<Self> {
            let device = Device::Cpu;

            let repo = Repo::with_revision(
                MODEL_REPO.to_string(),
                RepoType::Model,
                "main".to_string(),
            );
            let api = Api::new().map_err(|e| MindCoreError::Embedding(format!("HF API init: {e}")))?;
            let api = api.repo(repo);

            let config_path = api
                .get("config.json")
                .map_err(|e| MindCoreError::ModelNotAvailable(format!("config.json: {e}")))?;
            let tokenizer_path = api
                .get("tokenizer.json")
                .map_err(|e| MindCoreError::ModelNotAvailable(format!("tokenizer.json: {e}")))?;
            let weights_path = api
                .get("model.safetensors")
                .map_err(|e| MindCoreError::ModelNotAvailable(format!("model.safetensors: {e}")))?;

            Self::from_paths(&config_path, &tokenizer_path, &weights_path, device)
        }

        /// Create from pre-downloaded model files.
        ///
        /// Expected files in `model_dir`: `config.json`, `tokenizer.json`, `model.safetensors`
        pub fn from_path(model_dir: impl AsRef<Path>) -> Result<Self> {
            let dir = model_dir.as_ref();
            let config_path = dir.join("config.json");
            let tokenizer_path = dir.join("tokenizer.json");
            let weights_path = dir.join("model.safetensors");

            for path in [&config_path, &tokenizer_path, &weights_path] {
                if !path.exists() {
                    return Err(MindCoreError::ModelNotAvailable(format!(
                        "missing model file: {}",
                        path.display()
                    )));
                }
            }

            Self::from_paths(&config_path, &tokenizer_path, &weights_path, Device::Cpu)
        }

        /// Set Matryoshka dimension override (truncate vectors after embedding).
        pub fn with_dimensions_override(mut self, dims: usize) -> Self {
            self.dimensions_override = Some(dims);
            self
        }

        fn from_paths(
            config_path: &Path,
            tokenizer_path: &Path,
            weights_path: &Path,
            device: Device,
        ) -> Result<Self> {
            let config_str = std::fs::read_to_string(config_path)?;
            let config: Config = serde_json::from_str(&config_str)
                .map_err(|e| MindCoreError::Embedding(format!("config parse: {e}")))?;

            let tokenizer = Tokenizer::from_file(tokenizer_path)
                .map_err(|e| MindCoreError::Embedding(format!("tokenizer load: {e}")))?;

            // Safety: mmap is the standard way to load large model files.
            // The file must not be modified while mapped.
            #[allow(unsafe_code)]
            let vb = unsafe {
                VarBuilder::from_mmaped_safetensors(&[weights_path], DType::F32, &device)
                    .map_err(|e| MindCoreError::Embedding(format!("weights load: {e}")))?
            };

            // granite-small-r2 uses sentence-transformers naming (no "model." prefix)
            // but candle's ModernBert expects HuggingFace transformers naming.
            // Remap: when candle asks for "model.X.weight", look for "X.weight" in the file.
            let vb = vb.rename_f(|name| {
                name.strip_prefix("model.").unwrap_or(name).to_string()
            });

            let model = ModernBert::load(vb, &config)
                .map_err(|e| MindCoreError::Embedding(format!("model load: {e}")))?;

            Ok(Self {
                model,
                tokenizer,
                device,
                dimensions_override: None,
            })
        }

        /// Mean pooling with attention mask weighting.
        fn mean_pool(hidden_states: &Tensor, attention_mask: &Tensor) -> std::result::Result<Tensor, candle_core::Error> {
            let mask = attention_mask.to_dtype(DType::F32)?.unsqueeze(2)?;
            let sum_embeddings = hidden_states.broadcast_mul(&mask)?.sum(1)?;
            let sum_mask = mask.sum(1)?;
            sum_embeddings.broadcast_div(&sum_mask)
        }

        /// L2 normalization.
        fn normalize(v: &Tensor) -> std::result::Result<Tensor, candle_core::Error> {
            v.broadcast_div(&v.sqr()?.sum_keepdim(1)?.sqrt()?)
        }
    }

    impl EmbeddingBackend for CandleNativeBackend {
        fn embed(&self, text: &str) -> Result<Vec<f32>> {
            let results = self.embed_batch(&[text])?;
            results
                .into_iter()
                .next()
                .ok_or_else(|| MindCoreError::Embedding("empty batch result".into()))
        }

        fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
            let mut tokenizer = self.tokenizer.clone();
            tokenizer
                .with_padding(Some(PaddingParams {
                    strategy: PaddingStrategy::BatchLongest,
                    ..Default::default()
                }));

            let encodings = tokenizer
                .encode_batch(texts.to_vec(), true)
                .map_err(|e| MindCoreError::Embedding(format!("tokenize: {e}")))?;

            let token_ids: Vec<Tensor> = encodings
                .iter()
                .map(|enc| {
                    Tensor::new(enc.get_ids(), &self.device)
                        .map_err(|e| MindCoreError::Embedding(format!("tensor: {e}")))
                })
                .collect::<Result<Vec<_>>>()?;

            let attention_masks: Vec<Tensor> = encodings
                .iter()
                .map(|enc| {
                    Tensor::new(enc.get_attention_mask(), &self.device)
                        .map_err(|e| MindCoreError::Embedding(format!("mask tensor: {e}")))
                })
                .collect::<Result<Vec<_>>>()?;

            let token_ids = Tensor::stack(&token_ids, 0)
                .map_err(|e| MindCoreError::Embedding(format!("stack ids: {e}")))?;
            let attention_mask = Tensor::stack(&attention_masks, 0)
                .map_err(|e| MindCoreError::Embedding(format!("stack masks: {e}")))?;

            let hidden_states = self
                .model
                .forward(&token_ids, &attention_mask)
                .map_err(|e| MindCoreError::Embedding(format!("forward: {e}")))?;

            let pooled = Self::mean_pool(&hidden_states, &attention_mask)
                .map_err(|e| MindCoreError::Embedding(format!("pool: {e}")))?;

            let normalized = Self::normalize(&pooled)
                .map_err(|e| MindCoreError::Embedding(format!("normalize: {e}")))?;

            let batch_size = texts.len();
            let mut results = Vec::with_capacity(batch_size);
            for i in 0..batch_size {
                let mut vec: Vec<f32> = normalized
                    .get(i)
                    .map_err(|e| MindCoreError::Embedding(format!("get vec: {e}")))?
                    .to_vec1::<f32>()
                    .map_err(|e| MindCoreError::Embedding(format!("to_vec1: {e}")))?;

                // Matryoshka truncation
                if let Some(dims) = self.dimensions_override {
                    if dims < vec.len() {
                        vec.truncate(dims);
                        // Re-normalize after truncation
                        crate::embeddings::pooling::normalize_l2_inplace(&mut vec);
                    }
                }

                results.push(vec);
            }

            Ok(results)
        }

        fn dimensions(&self) -> usize {
            self.dimensions_override.unwrap_or(DIMENSIONS)
        }

        fn is_available(&self) -> bool {
            true
        }

        fn model_name(&self) -> &str {
            MODEL_NAME
        }
    }
}

#[cfg(feature = "local-embeddings")]
pub use inner::CandleNativeBackend;