gaze_recognizers/ner/
detector.rs

1use std::fmt;
2use std::path::{Path, PathBuf};
3use std::sync::Arc;
4
5use gaze_types::{Detection, Detector, RecognizerRuntimeError};
6
7use super::backend::{load_backend, NerBackend};
8use super::decode;
9use super::error::NerLoadError;
10use super::loader::warn_on_label_vocab_mismatch;
11use super::types::{LabelMap, NerBackendKind, NerOptions, NerSpanResult};
12
13/// NER detector backed by a pinned local model artifact set. Multiple
14/// `NerDetector` instances with different backends may be stacked in the
15/// same `Pipeline`; span-conflict resolution picks winners across detectors.
16pub struct NerDetector {
17    #[allow(dead_code)]
18    pub(crate) model_dir: PathBuf,
19    pub(crate) backend_kind: NerBackendKind,
20    pub(crate) recognizer_version_id: String,
21    pub(crate) locale: Option<String>,
22    pub(crate) threshold: f32,
23    pub(crate) backend: Arc<dyn NerBackend>,
24}
25
26impl fmt::Debug for NerDetector {
27    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
28        f.debug_struct("NerDetector")
29            .field("model_dir", &self.model_dir)
30            .field("backend_kind", &self.backend_kind)
31            .field("recognizer_version_id", &self.recognizer_version_id)
32            .field("locale", &self.locale)
33            .field("threshold", &self.threshold)
34            .finish_non_exhaustive()
35    }
36}
37
38impl NerDetector {
39    /// Full load: verify artifacts, initialize the configured backend.
40    /// Fails closed on any load error.
41    pub fn load(model_dir: &Path) -> Result<Self, NerLoadError> {
42        Self::load_with_options(model_dir, NerOptions::default())
43    }
44
45    pub fn load_with_options(model_dir: &Path, options: NerOptions) -> Result<Self, NerLoadError> {
46        let verified = Self::verify_artifacts(model_dir)?;
47        let backend_kind = verified.backend_kind;
48        let recognizer_version_id = format!(
49            "ner.{}.{}",
50            verified.recognizer_model_id, verified.recognizer_model_version
51        );
52        let model_dir_path = verified.model_dir.clone();
53        let label_count = verified.labels.len();
54        let id2label_len = verified.id2label.len();
55        warn_on_label_vocab_mismatch(&verified.labels, &verified.id2label, model_dir);
56        let backend = load_backend(verified)?;
57
58        tracing::info!(
59            backend = backend_kind.as_str(),
60            recognizer_version_id = %recognizer_version_id,
61            labels = label_count,
62            id2label_size = id2label_len,
63            locale = options.locale.as_deref().unwrap_or(""),
64            threshold = options.threshold,
65            model_dir = %model_dir_path.display(),
66            "ner: detector registered"
67        );
68
69        Ok(Self {
70            model_dir: model_dir_path,
71            backend_kind,
72            recognizer_version_id,
73            locale: options.locale,
74            threshold: options.threshold,
75            backend,
76        })
77    }
78
79    pub fn locale(&self) -> Option<&str> {
80        self.locale.as_deref()
81    }
82
83    pub fn backend_kind(&self) -> NerBackendKind {
84        self.backend_kind
85    }
86
87    pub fn recognizer_version_id(&self) -> &str {
88        &self.recognizer_version_id
89    }
90
91    pub(crate) fn detect_span_results(
92        &self,
93        input: &str,
94    ) -> Result<Vec<NerSpanResult>, super::error::NerRuntimeError> {
95        let mut spans = Vec::new();
96        for chunk in input_chunks(input) {
97            spans.extend(self.backend.detect(&input[chunk.clone()])?.into_iter().map(
98                |mut span| {
99                    span.span = span.span.start + chunk.start..span.span.end + chunk.start;
100                    span
101                },
102            ));
103        }
104        Ok(merge_overlapping_spans(spans))
105    }
106
107    /// Label/offset reconstruction helper. Public for testing the BIO merge.
108    /// `subword_spans` are byte ranges against the tokenizer input string,
109    /// `subword_labels` are CoNLL-style labels per subword (e.g. `O`, `B-PER`,
110    /// `I-PER`). Returns merged detections, dropping labels absent from the
111    /// label map and subword spans overlapping special tokens (empty ranges).
112    pub fn merge_bio_spans(
113        labels: &LabelMap,
114        subword_spans: &[(usize, usize)],
115        subword_labels: &[&str],
116        source: &str,
117    ) -> Vec<Detection> {
118        decode::merge_bio_spans(labels, subword_spans, subword_labels, source)
119    }
120
121    pub fn merge_bio_span_results(
122        labels: &LabelMap,
123        subword_spans: &[(usize, usize)],
124        subword_labels: &[&str],
125        subword_scores: &[f32],
126        source: &str,
127    ) -> Vec<NerSpanResult> {
128        decode::merge_bio_span_results(
129            labels,
130            subword_spans,
131            subword_labels,
132            subword_scores,
133            source,
134        )
135    }
136}
137
138impl Detector for NerDetector {
139    fn detect(&self, input: &str) -> Vec<Detection> {
140        self.try_detect(input)
141            .expect("ner detector backend failure is fail-closed")
142    }
143
144    fn try_detect(&self, input: &str) -> Result<Vec<Detection>, RecognizerRuntimeError> {
145        self.detect_span_results(input)
146            .map(|detections| {
147                detections
148                    .into_iter()
149                    .map(|span| {
150                        Detection::new(
151                            span.span,
152                            span.class,
153                            format!("ner/{}", self.backend_kind.as_str()),
154                        )
155                    })
156                    .collect()
157            })
158            .map_err(|err| {
159                tracing::warn!(backend = self.backend_kind.as_str(), error = %err, "ner: backend detect failed");
160                RecognizerRuntimeError::new("ner", err.to_string())
161            })
162    }
163}
164
165const NER_CHUNK_TOKEN_WINDOW: usize = 512;
166const NER_CHUNK_TOKEN_OVERLAP: usize = 32;
167
168fn input_chunks(input: &str) -> Vec<std::ops::Range<usize>> {
169    let mut tokens = Vec::new();
170    let mut token_start = None;
171    for (idx, ch) in input.char_indices() {
172        if ch.is_whitespace() {
173            if let Some(start) = token_start.take() {
174                tokens.push(start..idx);
175            }
176        } else if token_start.is_none() {
177            token_start = Some(idx);
178        }
179    }
180    if let Some(start) = token_start {
181        tokens.push(start..input.len());
182    }
183
184    if tokens.len() <= NER_CHUNK_TOKEN_WINDOW {
185        return std::iter::once(0..input.len()).collect();
186    }
187
188    let stride = NER_CHUNK_TOKEN_WINDOW - NER_CHUNK_TOKEN_OVERLAP;
189    let mut chunks = Vec::new();
190    let mut token_start = 0;
191    while token_start < tokens.len() {
192        let token_end = (token_start + NER_CHUNK_TOKEN_WINDOW).min(tokens.len());
193        chunks.push(tokens[token_start].start..tokens[token_end - 1].end);
194        if token_end == tokens.len() {
195            break;
196        }
197        token_start += stride;
198    }
199    chunks
200}
201
202fn merge_overlapping_spans(mut spans: Vec<NerSpanResult>) -> Vec<NerSpanResult> {
203    spans.sort_by(|left, right| {
204        left.span
205            .start
206            .cmp(&right.span.start)
207            .then(left.span.end.cmp(&right.span.end))
208            .then(left.class.cmp(&right.class))
209    });
210
211    let mut merged: Vec<NerSpanResult> = Vec::new();
212    for span in spans {
213        if let Some(last) = merged.last_mut() {
214            if last.class == span.class && last.span.end >= span.span.start {
215                last.span.end = last.span.end.max(span.span.end);
216                last.score = last.score.max(span.score);
217                continue;
218            }
219        }
220        merged.push(span);
221    }
222    merged
223}
gaze_recognizers/ner/detector.rs

gaze_recognizers/ner/
detector.rs
