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, CONFIG_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 labels = parse_labels(&model_dir.join(LABELS_FILE))?;
58        let config = parse_config(&model_dir.join(CONFIG_FILE))?;
59        let backend_kind = NerBackendKind::parse(config.backend.as_deref())?;
60        let id2label = config_to_id2label(config.id2label)?;
61
62        Ok(VerifiedArtifacts {
63            model_dir: model_dir.to_path_buf(),
64            backend_kind,
65            labels,
66            id2label,
67        })
68    }
69}
70
71pub(crate) fn warn_on_label_vocab_mismatch(
72    labels: &LabelMap,
73    id2label: &[String],
74    model_dir: &Path,
75) {
76    let mut usable = 0usize;
77    for tag in id2label {
78        let (_, entity) = split_bio(tag);
79        if entity.is_empty() {
80            continue;
81        }
82        if labels.resolve(tag, entity).is_some() {
83            usable += 1;
84        }
85    }
86    if usable == 0 {
87        let sample_label: String = labels.keys().take(5).collect::<Vec<_>>().join(",");
88        let sample_id: String = id2label
89            .iter()
90            .take(5)
91            .cloned()
92            .collect::<Vec<_>>()
93            .join(",");
94        tracing::warn!(
95            model_dir = %model_dir.display(),
96            label_keys = %sample_label,
97            id2label_sample = %sample_id,
98            "ner: labels.json has zero overlap with model id2label — detector will emit zero detections. Expected keys like 'PER'/'LOC' or 'B-PER'/'I-PER'."
99        );
100    }
101}
102
103fn parse_checksums(path: &Path) -> Result<Vec<(String, String)>, NerLoadError> {
104    let file = fs::File::open(path).map_err(|source| NerLoadError::Io {
105        path: path.to_path_buf(),
106        source,
107    })?;
108    let reader = BufReader::new(file);
109    let mut entries = Vec::new();
110    for (idx, line) in reader.lines().enumerate() {
111        let line_no = idx + 1;
112        let line = line.map_err(|source| NerLoadError::Io {
113            path: path.to_path_buf(),
114            source,
115        })?;
116        let trimmed = line.trim();
117        if trimmed.is_empty() || trimmed.starts_with('#') {
118            continue;
119        }
120        let mut parts = trimmed.splitn(2, char::is_whitespace);
121        let Some(hash) = parts.next() else {
122            return Err(NerLoadError::ChecksumsMalformed {
123                line: line_no,
124                reason: "missing hash".into(),
125            });
126        };
127        let Some(rest) = parts.next() else {
128            return Err(NerLoadError::ChecksumsMalformed {
129                line: line_no,
130                reason: "missing path".into(),
131            });
132        };
133        if hash.len() != 64 || !hash.chars().all(|c| c.is_ascii_hexdigit()) {
134            return Err(NerLoadError::ChecksumsMalformed {
135                line: line_no,
136                reason: format!("invalid sha256 hex: {hash}"),
137            });
138        }
139        let file = rest.trim_start().trim_start_matches('*').trim().to_string();
140        if file.is_empty() {
141            return Err(NerLoadError::ChecksumsMalformed {
142                line: line_no,
143                reason: "empty path".into(),
144            });
145        }
146        entries.push((file, hash.to_ascii_lowercase()));
147    }
148    Ok(entries)
149}
150
151fn hash_file(path: &Path) -> Result<String, NerLoadError> {
152    let bytes = fs::read(path).map_err(|source| NerLoadError::Io {
153        path: path.to_path_buf(),
154        source,
155    })?;
156    let mut hasher = Sha256::new();
157    hasher.update(&bytes);
158    Ok(hex::encode(hasher.finalize()))
159}
160
161fn parse_labels(path: &Path) -> Result<LabelMap, NerLoadError> {
162    let bytes = fs::read(path).map_err(|source| NerLoadError::Io {
163        path: path.to_path_buf(),
164        source,
165    })?;
166    let raw: BTreeMap<String, String> =
167        serde_json::from_slice(&bytes).map_err(|err| NerLoadError::LabelsParse(err.to_string()))?;
168    let mut map = BTreeMap::new();
169    for (key, value) in raw {
170        let class = match value.to_ascii_lowercase().as_str() {
171            "name" | "per" | "person" => PiiClass::Name,
172            "location" | "loc" => PiiClass::Location,
173            "organization" | "org" => PiiClass::Organization,
174            "email" => PiiClass::Email,
175            "drop" | "ignore" | "" => continue,
176            other => PiiClass::custom(other),
177        };
178        map.insert(key, class);
179    }
180    if map.is_empty() {
181        return Err(NerLoadError::LabelsParse(
182            "labels.json produced no usable mappings".into(),
183        ));
184    }
185    Ok(LabelMap(map))
186}
187
188#[derive(Deserialize)]
189struct ConfigFile {
190    backend: Option<String>,
191    id2label: Option<BTreeMap<String, String>>,
192}
193
194fn parse_config(path: &Path) -> Result<ConfigFile, NerLoadError> {
195    let bytes = fs::read(path).map_err(|source| NerLoadError::Io {
196        path: path.to_path_buf(),
197        source,
198    })?;
199    serde_json::from_slice(&bytes).map_err(|err| NerLoadError::ConfigParse(err.to_string()))
200}
201
202fn config_to_id2label(
203    id2label: Option<BTreeMap<String, String>>,
204) -> Result<Vec<String>, NerLoadError> {
205    let map = id2label
206        .ok_or_else(|| NerLoadError::ConfigParse("config.json missing id2label".to_string()))?;
207    let mut pairs: Vec<(usize, String)> = map
208        .into_iter()
209        .map(|(key, value)| {
210            key.parse::<usize>()
211                .map(|index| (index, value))
212                .map_err(|err| NerLoadError::ConfigParse(format!("id2label key {key}: {err}")))
213        })
214        .collect::<Result<_, _>>()?;
215    pairs.sort_by_key(|(index, _)| *index);
216    let max_idx = pairs.last().map(|(index, _)| *index).unwrap_or(0);
217    let mut out = vec!["O".to_string(); max_idx + 1];
218    for (index, label) in pairs {
219        out[index] = label;
220    }
221    Ok(out)
222}
gaze_recognizers/ner/loader.rs

gaze_recognizers/ner/
loader.rs
