hermes_core/tokenizer/
idf_weights.rs

1//! Pre-computed IDF weights from model's `idf.json`
2//!
3//! Neural sparse models (e.g., opensearch-neural-sparse-encoding-multilingual-v1)
4//! ship `idf.json` with IDF values calibrated during training. Using these weights
5//! instead of index-derived IDF produces correct rankings for doc-only models.
6//!
7//! The idf.json format maps **token strings** to IDF values:
8//! `{"[PAD]": 0.607, "hemoglobin": 8.12, "##ing": 1.05, ...}`
9//!
10//! At load time, we resolve token strings to numeric IDs via the model's tokenizer,
11//! then store as a flat `Vec<f32>` for O(1) lookup by token_id.
12
13#[cfg(feature = "native")]
14use std::collections::HashMap;
15#[cfg(feature = "native")]
16use std::path::Path;
17#[cfg(feature = "native")]
18use std::sync::Arc;
19
20#[cfg(feature = "native")]
21use log::{debug, warn};
22#[cfg(feature = "native")]
23use parking_lot::RwLock;
24
25#[cfg(feature = "native")]
26use crate::Result;
27#[cfg(feature = "native")]
28use crate::error::Error;
29
30/// Pre-computed IDF weights indexed by token_id
31///
32/// Stored as a flat `Vec<f32>` for O(1) lookup by token_id.
33/// For mBERT's 105K vocab this uses ~420KB of memory.
34#[cfg(feature = "native")]
35pub struct IdfWeights {
36    weights: Vec<f32>,
37}
38
39#[cfg(feature = "native")]
40impl IdfWeights {
41    /// Get the IDF weight for a token_id
42    ///
43    /// Returns 1.0 for out-of-range token_ids (neutral weight).
44    #[inline]
45    pub fn get(&self, token_id: u32) -> f32 {
46        self.weights.get(token_id as usize).copied().unwrap_or(1.0)
47    }
48
49    /// Load IDF weights from a JSON object, resolving token strings to IDs
50    /// via the provided tokenizer.
51    ///
52    /// The idf.json maps token strings → IDF values. We use `token_to_id`
53    /// to convert each key to a numeric token ID for O(1) lookup.
54    fn from_json_with_tokenizer(
55        json_bytes: &[u8],
56        tokenizer: &tokenizers::Tokenizer,
57    ) -> Result<Self> {
58        let map: HashMap<String, f64> = serde_json::from_slice(json_bytes)
59            .map_err(|e| Error::Tokenizer(format!("Failed to parse idf.json: {}", e)))?;
60
61        if map.is_empty() {
62            return Err(Error::Tokenizer("idf.json is empty".to_string()));
63        }
64
65        // Resolve token strings to IDs and find max ID
66        let mut resolved: Vec<(u32, f32)> = Vec::with_capacity(map.len());
67        let mut missed = 0u32;
68        for (token_str, value) in &map {
69            if let Some(id) = tokenizer.token_to_id(token_str) {
70                resolved.push((id, *value as f32));
71            } else {
72                missed += 1;
73            }
74        }
75
76        if resolved.is_empty() {
77            return Err(Error::Tokenizer(
78                "idf.json: no tokens could be resolved to IDs via tokenizer".to_string(),
79            ));
80        }
81
82        let max_id = resolved.iter().map(|(id, _)| *id).max().unwrap();
83
84        // Initialize with 1.0 (neutral weight) for unmapped tokens
85        let mut weights = vec![1.0f32; (max_id + 1) as usize];
86        for &(id, value) in &resolved {
87            weights[id as usize] = value;
88        }
89
90        debug!(
91            "Loaded {} IDF weights via tokenizer (vec size: {}, unresolved: {})",
92            resolved.len(),
93            weights.len(),
94            missed,
95        );
96
97        Ok(Self { weights })
98    }
99}
100
101/// Global cache for IDF weights, keyed by model name.
102/// Caches both successful loads and failures to avoid repeated download attempts.
103#[cfg(feature = "native")]
104pub struct IdfWeightsCache {
105    cache: RwLock<HashMap<String, Option<Arc<IdfWeights>>>>,
106}
107
108#[cfg(feature = "native")]
109impl Default for IdfWeightsCache {
110    fn default() -> Self {
111        Self::new()
112    }
113}
114
115#[cfg(feature = "native")]
116impl IdfWeightsCache {
117    /// Create a new IDF weights cache
118    pub fn new() -> Self {
119        Self {
120            cache: RwLock::new(HashMap::new()),
121        }
122    }
123
124    /// Get or load IDF weights for a model
125    ///
126    /// Lookup order:
127    /// 1. In-memory cache
128    /// 2. Local file in `cache_dir` (e.g. index directory): `idf_<sanitized_model>.json`
129    /// 3. HuggingFace hub download (saved to `cache_dir` on success)
130    ///
131    /// Returns `None` if `idf.json` is not available (graceful fallback).
132    /// Both successes and failures are cached to avoid repeated attempts.
133    pub fn get_or_load(
134        &self,
135        model_name: &str,
136        cache_dir: Option<&Path>,
137    ) -> Option<Arc<IdfWeights>> {
138        // Check in-memory cache first (covers both success and cached failure)
139        {
140            let cache = self.cache.read();
141            if let Some(entry) = cache.get(model_name) {
142                return entry.as_ref().map(Arc::clone);
143            }
144        }
145
146        // Try local cache file, then HF hub
147        match self.load_with_local_cache(model_name, cache_dir) {
148            Ok(weights) => {
149                let weights = Arc::new(weights);
150                let mut cache = self.cache.write();
151                cache.insert(model_name.to_string(), Some(Arc::clone(&weights)));
152                Some(weights)
153            }
154            Err(e) => {
155                warn!(
156                    "Could not load idf.json for model '{}': {}. Falling back to index-derived IDF.",
157                    model_name, e
158                );
159                let mut cache = self.cache.write();
160                cache.insert(model_name.to_string(), None);
161                None
162            }
163        }
164    }
165
166    /// Sanitize model name for use as a filename component
167    fn sanitized_model_name(model_name: &str) -> String {
168        model_name.replace('/', "--")
169    }
170
171    /// Local cache filename for a model's idf.json
172    fn local_cache_path(cache_dir: &Path, model_name: &str) -> std::path::PathBuf {
173        cache_dir.join(format!(
174            "idf_{}.json",
175            Self::sanitized_model_name(model_name)
176        ))
177    }
178
179    /// Try loading from local cache file first, then fall back to HF hub download.
180    /// On successful HF download, saves a copy to the local cache directory.
181    fn load_with_local_cache(
182        &self,
183        model_name: &str,
184        cache_dir: Option<&Path>,
185    ) -> Result<IdfWeights> {
186        let tokenizer = super::tokenizer_cache().get_or_load(model_name)?;
187
188        // Try local cache first
189        if let Some(dir) = cache_dir {
190            let local_path = Self::local_cache_path(dir, model_name);
191            if local_path.exists() {
192                let json_bytes = std::fs::read(&local_path).map_err(|e| {
193                    Error::Tokenizer(format!(
194                        "Failed to read cached idf.json at {:?}: {}",
195                        local_path, e
196                    ))
197                })?;
198                debug!(
199                    "Loaded idf.json from local cache: {:?} for model '{}'",
200                    local_path, model_name
201                );
202                return IdfWeights::from_json_with_tokenizer(&json_bytes, &tokenizer.tokenizer);
203            }
204        }
205
206        // Download from HF hub
207        let json_bytes = self.download_idf_json(model_name)?;
208
209        // Save to local cache for next time
210        if let Some(dir) = cache_dir {
211            let local_path = Self::local_cache_path(dir, model_name);
212            if let Err(e) = std::fs::write(&local_path, &json_bytes) {
213                warn!(
214                    "Failed to cache idf.json to {:?}: {} (non-fatal)",
215                    local_path, e
216                );
217            } else {
218                debug!(
219                    "Cached idf.json to {:?} for model '{}'",
220                    local_path, model_name
221                );
222            }
223        }
224
225        IdfWeights::from_json_with_tokenizer(&json_bytes, &tokenizer.tokenizer)
226    }
227
228    /// Load idf.json bytes, preferring the HF hub's local cache (no network)
229    /// and only downloading if not already cached.
230    fn download_idf_json(&self, model_name: &str) -> Result<Vec<u8>> {
231        let client = hf_hub::HFClientSync::new()
232            .map_err(|e| Error::Tokenizer(format!("Failed to create HF hub client: {}", e)))?;
233        let (owner, name) = hf_hub::split_id(model_name);
234        let repo = client.model(owner, name);
235
236        // Check HF hub's disk cache first — no network request
237        if let Ok(cached_path) = repo
238            .download_file()
239            .filename("idf.json")
240            .local_files_only(true)
241            .send()
242        {
243            debug!(
244                "Loaded idf.json from HF cache: {:?} for model '{}'",
245                cached_path, model_name
246            );
247            return std::fs::read(&cached_path).map_err(|e| {
248                Error::Tokenizer(format!(
249                    "Failed to read cached idf.json at {:?}: {}",
250                    cached_path, e
251                ))
252            });
253        }
254
255        // Not in cache — download via API
256        let idf_path = repo
257            .download_file()
258            .filename("idf.json")
259            .send()
260            .map_err(|e| {
261                Error::Tokenizer(format!(
262                    "Failed to download idf.json from '{}': {}",
263                    model_name, e
264                ))
265            })?;
266
267        debug!(
268            "Downloaded idf.json from '{}' to {:?}",
269            model_name, idf_path
270        );
271
272        std::fs::read(&idf_path).map_err(|e| {
273            Error::Tokenizer(format!("Failed to read idf.json at {:?}: {}", idf_path, e))
274        })
275    }
276
277    /// Clear the cache
278    pub fn clear(&self) {
279        let mut cache = self.cache.write();
280        cache.clear();
281    }
282}
283
284/// Global IDF weights cache instance
285#[cfg(feature = "native")]
286static IDF_WEIGHTS_CACHE: std::sync::OnceLock<IdfWeightsCache> = std::sync::OnceLock::new();
287
288/// Get the global IDF weights cache
289#[cfg(feature = "native")]
290pub fn idf_weights_cache() -> &'static IdfWeightsCache {
291    IDF_WEIGHTS_CACHE.get_or_init(IdfWeightsCache::new)
292}
293
294#[cfg(test)]
295#[cfg(feature = "native")]
296mod tests {
297    use super::*;
298
299    /// Build a minimal tokenizer with a known vocab for testing
300    fn test_tokenizer() -> tokenizers::Tokenizer {
301        use tokenizers::models::wordpiece::WordPiece;
302        let wp = WordPiece::builder()
303            .vocab([
304                ("[UNK]".to_string(), 0),
305                ("hello".to_string(), 1),
306                ("world".to_string(), 2),
307                ("foo".to_string(), 5),
308                ("bar".to_string(), 100),
309            ])
310            .unk_token("[UNK]".into())
311            .build()
312            .unwrap();
313        tokenizers::Tokenizer::new(wp)
314    }
315
316    #[test]
317    fn test_idf_weights_from_json_with_tokenizer() {
318        let json = br#"{"hello": 1.5, "world": 2.0, "foo": 0.5, "bar": 3.0}"#;
319        let tokenizer = test_tokenizer();
320        let weights = IdfWeights::from_json_with_tokenizer(json, &tokenizer).unwrap();
321
322        // hello=1, world=2, foo=5, bar=100
323        assert!((weights.get(1) - 1.5).abs() < f32::EPSILON);
324        assert!((weights.get(2) - 2.0).abs() < f32::EPSILON);
325        assert!((weights.get(5) - 0.5).abs() < f32::EPSILON);
326        assert!((weights.get(100) - 3.0).abs() < f32::EPSILON);
327
328        // Unmapped tokens get 1.0
329        assert!((weights.get(3) - 1.0).abs() < f32::EPSILON);
330        assert!((weights.get(50) - 1.0).abs() < f32::EPSILON);
331
332        // Out-of-range tokens get 1.0
333        assert!((weights.get(999) - 1.0).abs() < f32::EPSILON);
334    }
335
336    #[test]
337    fn test_idf_weights_unresolvable_tokens_skipped() {
338        // "unknown_xyz" is not in the vocab, should be skipped
339        let json = br#"{"hello": 1.5, "unknown_xyz": 9.9}"#;
340        let tokenizer = test_tokenizer();
341        let weights = IdfWeights::from_json_with_tokenizer(json, &tokenizer).unwrap();
342
343        assert!((weights.get(1) - 1.5).abs() < f32::EPSILON); // hello resolved
344    }
345
346    #[test]
347    fn test_idf_weights_empty_json() {
348        let json = br#"{}"#;
349        let tokenizer = test_tokenizer();
350        assert!(IdfWeights::from_json_with_tokenizer(json, &tokenizer).is_err());
351    }
352
353    #[test]
354    fn test_idf_weights_invalid_json() {
355        let json = br#"not json"#;
356        let tokenizer = test_tokenizer();
357        assert!(IdfWeights::from_json_with_tokenizer(json, &tokenizer).is_err());
358    }
359
360    #[test]
361    fn test_idf_weights_cache_structure() {
362        let cache = IdfWeightsCache::new();
363        assert!(cache.cache.read().is_empty());
364    }
365
366    #[test]
367    fn test_idf_weights_cache_miss_graceful() {
368        let cache = IdfWeightsCache::new();
369        // Non-existent model should return None gracefully
370        let result = cache.get_or_load("nonexistent-model-xyz-12345", None);
371        assert!(result.is_none());
372    }
373}
hermes_core/tokenizer/idf_weights.rs

hermes_core/tokenizer/
idf_weights.rs
