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) 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    /// Label/offset reconstruction helper. Public for testing the BIO merge.
92    /// `subword_spans` are byte ranges against the tokenizer input string,
93    /// `subword_labels` are CoNLL-style labels per subword (e.g. `O`, `B-PER`,
94    /// `I-PER`). Returns merged detections, dropping labels absent from the
95    /// label map and subword spans overlapping special tokens (empty ranges).
96    pub fn merge_bio_spans(
97        labels: &LabelMap,
98        subword_spans: &[(usize, usize)],
99        subword_labels: &[&str],
100        source: &str,
101    ) -> Vec<Detection> {
102        decode::merge_bio_spans(labels, subword_spans, subword_labels, source)
103    }
104
105    pub fn merge_bio_span_results(
106        labels: &LabelMap,
107        subword_spans: &[(usize, usize)],
108        subword_labels: &[&str],
109        subword_scores: &[f32],
110        source: &str,
111    ) -> Vec<NerSpanResult> {
112        decode::merge_bio_span_results(
113            labels,
114            subword_spans,
115            subword_labels,
116            subword_scores,
117            source,
118        )
119    }
120}
121
122impl Detector for NerDetector {
123    fn detect(&self, input: &str) -> Vec<Detection> {
124        match self.backend.detect(input) {
125            Ok(detections) => detections
126                .into_iter()
127                .map(|span| {
128                    Detection::new(
129                        span.span,
130                        span.class,
131                        format!("ner/{}", self.backend_kind.as_str()),
132                    )
133                })
134                .collect(),
135            Err(err) => {
136                tracing::warn!(backend = self.backend_kind.as_str(), error = %err, "ner: backend detect failed");
137                Vec::new()
138            }
139        }
140    }
141}
gaze_recognizers/ner/detector.rs

gaze_recognizers/ner/
detector.rs
