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 self.backend.chunk_ranges(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            .into_iter()
106            .filter(|span| decode::is_valid_entity_span(input, &span.span, &span.class, false))
107            .collect())
108    }
109
110    /// Label/offset reconstruction helper. Public for testing the BIO merge.
111    /// `subword_spans` are byte ranges against the tokenizer input string,
112    /// `subword_labels` are CoNLL-style labels per subword (e.g. `O`, `B-PER`,
113    /// `I-PER`). Returns merged detections, dropping labels absent from the
114    /// label map and subword spans overlapping special tokens (empty ranges).
115    pub fn merge_bio_spans(
116        labels: &LabelMap,
117        subword_spans: &[(usize, usize)],
118        subword_labels: &[&str],
119        source: &str,
120    ) -> Vec<Detection> {
121        decode::merge_bio_spans(labels, subword_spans, subword_labels, source)
122    }
123
124    pub fn merge_bio_span_results(
125        labels: &LabelMap,
126        subword_spans: &[(usize, usize)],
127        subword_labels: &[&str],
128        subword_scores: &[f32],
129        source: &str,
130    ) -> Vec<NerSpanResult> {
131        decode::merge_bio_span_results(
132            labels,
133            subword_spans,
134            subword_labels,
135            subword_scores,
136            source,
137        )
138    }
139}
140
141impl Detector for NerDetector {
142    fn detect(&self, input: &str) -> Vec<Detection> {
143        self.try_detect(input)
144            .expect("ner detector backend failure is fail-closed")
145    }
146
147    fn try_detect(&self, input: &str) -> Result<Vec<Detection>, RecognizerRuntimeError> {
148        self.detect_span_results(input)
149            .map(|detections| {
150                detections
151                    .into_iter()
152                    .map(|span| {
153                        Detection::new(
154                            span.span,
155                            span.class,
156                            format!("ner/{}", self.backend_kind.as_str()),
157                        )
158                    })
159                    .collect()
160            })
161            .map_err(|err| {
162                tracing::warn!(backend = self.backend_kind.as_str(), error = %err, "ner: backend detect failed");
163                RecognizerRuntimeError::new("ner", err.to_string())
164            })
165    }
166}
167
168fn merge_overlapping_spans(mut spans: Vec<NerSpanResult>) -> Vec<NerSpanResult> {
169    spans.sort_by(|left, right| {
170        left.span
171            .start
172            .cmp(&right.span.start)
173            .then(left.span.end.cmp(&right.span.end))
174            .then(left.class.cmp(&right.class))
175    });
176
177    let mut merged: Vec<NerSpanResult> = Vec::new();
178    for span in spans {
179        if let Some(last) = merged.last_mut() {
180            if last.class == span.class && last.span.end >= span.span.start {
181                last.span.end = last.span.end.max(span.span.end);
182                last.score = last.score.max(span.score);
183                continue;
184            }
185        }
186        merged.push(span);
187    }
188    merged
189}
gaze_recognizers/ner/detector.rs

gaze_recognizers/ner/
detector.rs
