gaze_recognizers/ner/
loader.rs

1use std::collections::BTreeMap;
2use std::fs;
3use std::io::{BufRead, BufReader};
4use std::path::Path;
5
6use gaze_types::PiiClass;
7use serde::Deserialize;
8use sha2::{Digest, Sha256};
9
10use super::decode::split_bio;
11use super::detector::NerDetector;
12use super::error::NerLoadError;
13use super::types::{
14    LabelMap, NerBackendKind, VerifiedArtifacts, CHECKSUMS_FILE, CONFIG_FILE, LABELS_FILE,
15    MODEL_FILE, TOKENIZER_FILE,
16};
17
18impl NerDetector {
19    /// Verify SHA256SUMS, parse labels.json and config.json. No backend
20    /// initialization. Used both by `load` and by unit tests.
21    pub fn verify_artifacts(model_dir: &Path) -> Result<VerifiedArtifacts, NerLoadError> {
22        if !model_dir.is_dir() {
23            return Err(NerLoadError::ModelDirMissing {
24                path: model_dir.to_path_buf(),
25            });
26        }
27
28        let sums_path = model_dir.join(CHECKSUMS_FILE);
29        if !sums_path.exists() {
30            return Err(NerLoadError::ChecksumsMissing { path: sums_path });
31        }
32
33        let entries = parse_checksums(&sums_path)?;
34        for required in [MODEL_FILE, TOKENIZER_FILE, LABELS_FILE] {
35            if !entries.iter().any(|(name, _)| name == required) {
36                return Err(NerLoadError::MissingArtifact {
37                    path: model_dir.join(required),
38                });
39            }
40        }
41
42        for (rel, expected) in &entries {
43            let path = model_dir.join(rel);
44            if !path.exists() {
45                return Err(NerLoadError::MissingArtifact { path });
46            }
47            let got = hash_file(&path)?;
48            if !got.eq_ignore_ascii_case(expected) {
49                return Err(NerLoadError::ChecksumMismatch {
50                    path,
51                    expected: expected.clone(),
52                    got,
53                });
54            }
55        }
56
57        let parsed_labels = parse_labels(&model_dir.join(LABELS_FILE))?;
58        let config = if entries.iter().any(|(name, _)| name == CONFIG_FILE) {
59            parse_config(&model_dir.join(CONFIG_FILE))?
60        } else if let Some(kiji_config) = parsed_labels.kiji_config {
61            kiji_config
62        } else {
63            return Err(NerLoadError::MissingArtifact {
64                path: model_dir.join(CONFIG_FILE),
65            });
66        };
67        let backend_kind = NerBackendKind::parse(config.backend.as_deref())?;
68        let recognizer_model_id = config
69            .recognizer_model_id()
70            .map(sanitize_recognizer_segment)
71            .filter(|value| !value.is_empty())
72            .unwrap_or_else(|| "unknown".to_string());
73        let recognizer_model_version = config
74            .recognizer_model_version()
75            .map(sanitize_version_segment)
76            .filter(|value| !value.is_empty())
77            .unwrap_or_else(|| "v0".to_string());
78        let id2label = config_to_id2label(config.id2label)?;
79
80        Ok(VerifiedArtifacts {
81            model_dir: model_dir.to_path_buf(),
82            backend_kind,
83            recognizer_model_id,
84            recognizer_model_version,
85            labels: parsed_labels.labels,
86            id2label,
87        })
88    }
89}
90
91pub(crate) fn warn_on_label_vocab_mismatch(
92    labels: &LabelMap,
93    id2label: &[String],
94    model_dir: &Path,
95) {
96    let mut usable = 0usize;
97    for tag in id2label {
98        let (_, entity) = split_bio(tag);
99        if entity.is_empty() {
100            continue;
101        }
102        if labels.resolve(tag, entity).is_some() {
103            usable += 1;
104        }
105    }
106    if usable == 0 {
107        let sample_label: String = labels.keys().take(5).collect::<Vec<_>>().join(",");
108        let sample_id: String = id2label
109            .iter()
110            .take(5)
111            .cloned()
112            .collect::<Vec<_>>()
113            .join(",");
114        tracing::warn!(
115            model_dir = %model_dir.display(),
116            label_keys = %sample_label,
117            id2label_sample = %sample_id,
118            "ner: labels.json has zero overlap with model id2label — detector will emit zero detections. Expected keys like 'PER'/'LOC' or 'B-PER'/'I-PER'."
119        );
120    }
121}
122
123fn parse_checksums(path: &Path) -> Result<Vec<(String, String)>, NerLoadError> {
124    let file = fs::File::open(path).map_err(|source| NerLoadError::Io {
125        path: path.to_path_buf(),
126        source,
127    })?;
128    let reader = BufReader::new(file);
129    let mut entries = Vec::new();
130    for (idx, line) in reader.lines().enumerate() {
131        let line_no = idx + 1;
132        let line = line.map_err(|source| NerLoadError::Io {
133            path: path.to_path_buf(),
134            source,
135        })?;
136        let trimmed = line.trim();
137        if trimmed.is_empty() || trimmed.starts_with('#') {
138            continue;
139        }
140        let mut parts = trimmed.splitn(2, char::is_whitespace);
141        let Some(hash) = parts.next() else {
142            return Err(NerLoadError::ChecksumsMalformed {
143                line: line_no,
144                reason: "missing hash".into(),
145            });
146        };
147        let Some(rest) = parts.next() else {
148            return Err(NerLoadError::ChecksumsMalformed {
149                line: line_no,
150                reason: "missing path".into(),
151            });
152        };
153        if hash.len() != 64 || !hash.chars().all(|c| c.is_ascii_hexdigit()) {
154            return Err(NerLoadError::ChecksumsMalformed {
155                line: line_no,
156                reason: format!("invalid sha256 hex: {hash}"),
157            });
158        }
159        let file = rest.trim_start().trim_start_matches('*').trim().to_string();
160        if file.is_empty() {
161            return Err(NerLoadError::ChecksumsMalformed {
162                line: line_no,
163                reason: "empty path".into(),
164            });
165        }
166        entries.push((file, hash.to_ascii_lowercase()));
167    }
168    Ok(entries)
169}
170
171fn hash_file(path: &Path) -> Result<String, NerLoadError> {
172    let bytes = fs::read(path).map_err(|source| NerLoadError::Io {
173        path: path.to_path_buf(),
174        source,
175    })?;
176    let mut hasher = Sha256::new();
177    hasher.update(&bytes);
178    Ok(hex::encode(hasher.finalize()))
179}
180
181struct ParsedLabels {
182    labels: LabelMap,
183    kiji_config: Option<ConfigFile>,
184}
185
186fn parse_labels(path: &Path) -> Result<ParsedLabels, NerLoadError> {
187    let bytes = fs::read(path).map_err(|source| NerLoadError::Io {
188        path: path.to_path_buf(),
189        source,
190    })?;
191    if let Ok(raw) = serde_json::from_slice::<BTreeMap<String, String>>(&bytes) {
192        return parse_flat_labels(raw).map(|labels| ParsedLabels {
193            labels,
194            kiji_config: None,
195        });
196    }
197
198    let raw: KijiLabelsFile =
199        serde_json::from_slice(&bytes).map_err(|err| NerLoadError::LabelsParse(err.to_string()))?;
200    parse_kiji_labels(raw)
201}
202
203fn parse_flat_labels(raw: BTreeMap<String, String>) -> Result<LabelMap, NerLoadError> {
204    let mut map = BTreeMap::new();
205    for (key, value) in raw {
206        let class = match value.to_ascii_lowercase().as_str() {
207            "name" | "per" | "person" => PiiClass::Name,
208            "location" | "loc" => PiiClass::Location,
209            "organization" | "org" => PiiClass::Organization,
210            "email" => PiiClass::Email,
211            "drop" | "ignore" | "" => continue,
212            other => PiiClass::custom(other),
213        };
214        map.insert(key, class);
215    }
216    if map.is_empty() {
217        return Err(NerLoadError::LabelsParse(
218            "labels.json produced no usable mappings".into(),
219        ));
220    }
221    Ok(LabelMap(map))
222}
223
224#[derive(Deserialize)]
225struct KijiLabelsFile {
226    schema_version: u64,
227    source: String,
228    source_commit: String,
229    labels: Vec<KijiLabelEntry>,
230}
231
232#[derive(Deserialize)]
233struct KijiLabelEntry {
234    id: String,
235    upstream: Vec<String>,
236}
237
238fn parse_kiji_labels(raw: KijiLabelsFile) -> Result<ParsedLabels, NerLoadError> {
239    if raw.schema_version != 1
240        || raw.source != "onnx-community/distilbert-NER-ONNX"
241        || raw.source_commit != "3a19fe9404a4469d91aa3d551558a97f68872f67"
242    {
243        return Err(NerLoadError::LabelsParse(
244            "unsupported Kiji labels.json manifest".to_string(),
245        ));
246    }
247
248    let mut labels = BTreeMap::new();
249    let mut id2label = vec!["O".to_string()];
250    for entry in raw.labels {
251        let class = match entry.id.as_str() {
252            "person" => PiiClass::Name,
253            "location" => PiiClass::Location,
254            "organization" => PiiClass::Organization,
255            "miscellaneous" => PiiClass::custom("miscellaneous"),
256            other => {
257                return Err(NerLoadError::LabelsParse(format!(
258                    "unsupported Kiji label id `{other}`"
259                )));
260            }
261        };
262        labels.insert(entry.id, class.clone());
263        for upstream in entry.upstream {
264            if upstream.is_empty() || upstream.contains('/') || upstream.contains('\\') {
265                return Err(NerLoadError::LabelsParse(
266                    "invalid Kiji upstream label".to_string(),
267                ));
268            }
269            if !id2label.iter().any(|existing| existing == &upstream) {
270                id2label.push(upstream.clone());
271            }
272            labels.insert(upstream, class.clone());
273        }
274    }
275    if labels.is_empty() {
276        return Err(NerLoadError::LabelsParse(
277            "labels.json produced no usable mappings".into(),
278        ));
279    }
280
281    let id2label = id2label
282        .into_iter()
283        .enumerate()
284        .map(|(index, label)| (index.to_string(), label))
285        .collect();
286    Ok(ParsedLabels {
287        labels: LabelMap(labels),
288        kiji_config: Some(ConfigFile {
289            backend: Some("ort".to_string()),
290            model_id: Some(raw.source),
291            model_name: None,
292            model: None,
293            version: None,
294            model_version: Some(raw.source_commit),
295            recognizer_version: None,
296            id2label: Some(id2label),
297        }),
298    })
299}
300
301#[derive(Deserialize)]
302struct ConfigFile {
303    backend: Option<String>,
304    model_id: Option<String>,
305    model_name: Option<String>,
306    model: Option<String>,
307    version: Option<String>,
308    model_version: Option<String>,
309    recognizer_version: Option<String>,
310    id2label: Option<BTreeMap<String, String>>,
311}
312
313impl ConfigFile {
314    fn recognizer_model_id(&self) -> Option<&str> {
315        self.model_id
316            .as_deref()
317            .or(self.model_name.as_deref())
318            .or(self.model.as_deref())
319    }
320
321    fn recognizer_model_version(&self) -> Option<&str> {
322        self.recognizer_version
323            .as_deref()
324            .or(self.model_version.as_deref())
325            .or(self.version.as_deref())
326    }
327}
328
329fn parse_config(path: &Path) -> Result<ConfigFile, NerLoadError> {
330    let bytes = fs::read(path).map_err(|source| NerLoadError::Io {
331        path: path.to_path_buf(),
332        source,
333    })?;
334    serde_json::from_slice(&bytes).map_err(|err| NerLoadError::ConfigParse(err.to_string()))
335}
336
337fn config_to_id2label(
338    id2label: Option<BTreeMap<String, String>>,
339) -> Result<Vec<String>, NerLoadError> {
340    let map = id2label
341        .ok_or_else(|| NerLoadError::ConfigParse("config.json missing id2label".to_string()))?;
342    let mut pairs: Vec<(usize, String)> = map
343        .into_iter()
344        .map(|(key, value)| {
345            key.parse::<usize>()
346                .map(|index| (index, value))
347                .map_err(|err| NerLoadError::ConfigParse(format!("id2label key {key}: {err}")))
348        })
349        .collect::<Result<_, _>>()?;
350    pairs.sort_by_key(|(index, _)| *index);
351    let max_idx = pairs.last().map(|(index, _)| *index).unwrap_or(0);
352    let mut out = vec!["O".to_string(); max_idx + 1];
353    for (index, label) in pairs {
354        out[index] = label;
355    }
356    Ok(out)
357}
358
359fn sanitize_recognizer_segment(raw: &str) -> String {
360    raw.trim()
361        .to_ascii_lowercase()
362        .chars()
363        .map(|ch| match ch {
364            'a'..='z' | '0'..='9' => ch,
365            _ => '-',
366        })
367        .collect::<String>()
368        .split('-')
369        .filter(|part| !part.is_empty())
370        .collect::<Vec<_>>()
371        .join("-")
372}
373
374fn sanitize_version_segment(raw: &str) -> String {
375    let sanitized = sanitize_recognizer_segment(raw);
376    if sanitized.is_empty() || sanitized.starts_with('v') {
377        sanitized
378    } else {
379        format!("v{sanitized}")
380    }
381}
gaze_recognizers/ner/loader.rs

gaze_recognizers/ner/
loader.rs
