gaze_recognizers/ner/
detector.rs

1use std::fmt;
2use std::path::{Path, PathBuf};
3use std::sync::Arc;
4
5use gaze_types::{Detection, Detector};
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) locale: Option<String>,
21    pub(crate) threshold: f32,
22    pub(crate) backend: Arc<dyn NerBackend>,
23}
24
25impl fmt::Debug for NerDetector {
26    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
27        f.debug_struct("NerDetector")
28            .field("model_dir", &self.model_dir)
29            .field("backend_kind", &self.backend_kind)
30            .field("locale", &self.locale)
31            .field("threshold", &self.threshold)
32            .finish_non_exhaustive()
33    }
34}
35
36impl NerDetector {
37    /// Full load: verify artifacts, initialize the configured backend.
38    /// Fails closed on any load error.
39    pub fn load(model_dir: &Path) -> Result<Self, NerLoadError> {
40        Self::load_with_options(model_dir, NerOptions::default())
41    }
42
43    pub fn load_with_options(model_dir: &Path, options: NerOptions) -> Result<Self, NerLoadError> {
44        let verified = Self::verify_artifacts(model_dir)?;
45        let backend_kind = verified.backend_kind;
46        let model_dir_path = verified.model_dir.clone();
47        let label_count = verified.labels.len();
48        let id2label_len = verified.id2label.len();
49        warn_on_label_vocab_mismatch(&verified.labels, &verified.id2label, model_dir);
50        let backend = load_backend(verified)?;
51
52        tracing::info!(
53            backend = backend_kind.as_str(),
54            labels = label_count,
55            id2label_size = id2label_len,
56            locale = options.locale.as_deref().unwrap_or(""),
57            threshold = options.threshold,
58            model_dir = %model_dir_path.display(),
59            "ner: detector registered"
60        );
61
62        Ok(Self {
63            model_dir: model_dir_path,
64            backend_kind,
65            locale: options.locale,
66            threshold: options.threshold,
67            backend,
68        })
69    }
70
71    pub fn locale(&self) -> Option<&str> {
72        self.locale.as_deref()
73    }
74
75    pub fn backend_kind(&self) -> NerBackendKind {
76        self.backend_kind
77    }
78
79    /// Label/offset reconstruction helper. Public for testing the BIO merge.
80    /// `subword_spans` are byte ranges against the tokenizer input string,
81    /// `subword_labels` are CoNLL-style labels per subword (e.g. `O`, `B-PER`,
82    /// `I-PER`). Returns merged detections, dropping labels absent from the
83    /// label map and subword spans overlapping special tokens (empty ranges).
84    pub fn merge_bio_spans(
85        labels: &LabelMap,
86        subword_spans: &[(usize, usize)],
87        subword_labels: &[&str],
88        source: &str,
89    ) -> Vec<Detection> {
90        decode::merge_bio_spans(labels, subword_spans, subword_labels, source)
91    }
92
93    pub fn merge_bio_span_results(
94        labels: &LabelMap,
95        subword_spans: &[(usize, usize)],
96        subword_labels: &[&str],
97        subword_scores: &[f32],
98        source: &str,
99    ) -> Vec<NerSpanResult> {
100        decode::merge_bio_span_results(
101            labels,
102            subword_spans,
103            subword_labels,
104            subword_scores,
105            source,
106        )
107    }
108}
109
110impl Detector for NerDetector {
111    fn detect(&self, input: &str) -> Vec<Detection> {
112        match self.backend.detect(input) {
113            Ok(detections) => detections
114                .into_iter()
115                .map(|span| {
116                    Detection::new(
117                        span.span,
118                        span.class,
119                        format!("ner/{}", self.backend_kind.as_str()),
120                    )
121                })
122                .collect(),
123            Err(err) => {
124                tracing::warn!(backend = self.backend_kind.as_str(), error = %err, "ner: backend detect failed");
125                Vec::new()
126            }
127        }
128    }
129}
gaze_recognizers/ner/detector.rs

gaze_recognizers/ner/
detector.rs
