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

gaze_recognizers/ner/
loader.rs
