gaze_recognizers/ner/
types.rs

1use std::collections::BTreeMap;
2use std::path::PathBuf;
3
4use gaze_types::PiiClass;
5
6use super::error::NerLoadError;
7
8/// Relative file names that must be present in a model directory.
9pub const MODEL_FILE: &str = "model.onnx";
10pub const TOKENIZER_FILE: &str = "tokenizer.json";
11pub const CONFIG_FILE: &str = "config.json";
12pub const LABELS_FILE: &str = "labels.json";
13pub const CHECKSUMS_FILE: &str = "SHA256SUMS";
14
15/// Labels file format.
16///
17/// Accepts two equivalent shapes so adopters aren't silently blocked by a
18/// keying convention mismatch:
19///
20/// - **Bare entity keys** (preferred, short): `{ "PER": "Name", "LOC": "Location" }`.
21/// - **BIO-prefixed keys** (mirrors CoNLL / HuggingFace `id2label` shape):
22///   `{ "B-PER": "Name", "I-PER": "Name", "B-LOC": "Location", "I-LOC": "Location" }`.
23///
24/// Values are matched against `PiiClass` variants by lowercase name, falling
25/// back to `PiiClass::custom(value)` if no built-in matches. The sentinel
26/// value `"drop"` (or `"ignore"`, `""`) removes the entry entirely so the
27/// detector silently skips that label.
28///
29/// Lookup (`resolve`) tries the full BIO tag first, then the bare entity
30/// type. Mixing both key shapes in a single file is allowed.
31#[derive(Debug, Clone, PartialEq, Eq)]
32pub struct LabelMap(pub BTreeMap<String, PiiClass>);
33
34impl LabelMap {
35    pub fn get(&self, conll_label: &str) -> Option<&PiiClass> {
36        self.0.get(conll_label)
37    }
38
39    /// Resolve a CoNLL subword tag (e.g. `"B-PER"`, `"I-LOC"`, `"PER"`) to a
40    /// `PiiClass`. Accepts both BIO-prefixed labels.json entries and bare
41    /// entity-type entries; tries the full tag first, then the stripped
42    /// entity.
43    pub fn resolve(&self, tag: &str, entity: &str) -> Option<&PiiClass> {
44        self.0.get(tag).or_else(|| self.0.get(entity))
45    }
46
47    /// Number of retained mappings (after `"drop"`/`"ignore"` sentinels are
48    /// filtered out by the parser). Used by the NER bootstrap `tracing::info!`
49    /// so adopters can tell at a glance whether their labels.json is empty
50    /// or misaligned.
51    pub fn len(&self) -> usize {
52        self.0.len()
53    }
54
55    pub fn is_empty(&self) -> bool {
56        self.0.is_empty()
57    }
58
59    /// Iterate the retained label keys. Used at load time to warn when
60    /// labels.json has zero overlap with the model's `id2label` vocab — a
61    /// silent-no-op symptom adopters otherwise only catch via missing
62    /// detections at runtime.
63    pub fn keys(&self) -> impl Iterator<Item = &str> {
64        self.0.keys().map(String::as_str)
65    }
66}
67
68#[derive(Debug, Clone, PartialEq)]
69pub struct NerOptions {
70    pub locale: Option<String>,
71    pub threshold: f32,
72}
73
74impl Default for NerOptions {
75    fn default() -> Self {
76        Self {
77            locale: None,
78            threshold: 0.3,
79        }
80    }
81}
82
83#[derive(Debug, Clone, PartialEq)]
84pub struct NerSpanResult {
85    pub span: std::ops::Range<usize>,
86    pub class: PiiClass,
87    pub score: f32,
88}
89
90/// Driver-style enum for NER backends. Backends are swappable under a common
91/// `NerBackend` trait; each owns its own model-specific state. Multiple
92/// `NerDetector` instances (e.g. a BERT token-classifier plus a GLiNER
93/// zero-shot model) can be stacked in the same `Pipeline` — span-conflict
94/// resolution picks winners across all detectors.
95#[derive(Debug, Clone, Copy, PartialEq, Eq)]
96pub enum NerBackendKind {
97    /// Standard BERT-family token classifier: fixed label vocabulary, BIO/IOB2
98    /// subword tagging, merged via `merge_bio_spans`. Driven by ONNX Runtime.
99    Ort,
100    /// GLiNER-family zero-shot / open-schema extractor: entity type strings
101    /// passed at inference, output is a span-score matrix. Shape reserved;
102    /// backend implementation lands when GLiNER artifacts are pinned.
103    Gliner,
104}
105
106impl NerBackendKind {
107    pub(crate) fn parse(raw: Option<&str>) -> Result<Self, NerLoadError> {
108        match raw.map(str::trim).filter(|value| !value.is_empty()) {
109            None => Ok(Self::Ort),
110            Some("ort") | Some("onnxruntime") | Some("bert-ort") => Ok(Self::Ort),
111            Some("gliner") | Some("gliner-ort") => Ok(Self::Gliner),
112            Some(other) => Err(NerLoadError::UnsupportedBackend {
113                backend: other.to_string(),
114            }),
115        }
116    }
117
118    pub(crate) fn as_str(self) -> &'static str {
119        match self {
120            Self::Ort => "ort",
121            Self::Gliner => "gliner",
122        }
123    }
124}
125
126/// Verified artifact handles. Produced by `verify_artifacts`, consumed by
127/// `NerDetector::load`. Split out so the load contract can be exercised by
128/// unit tests without initializing a backend runtime.
129#[derive(Debug, Clone)]
130pub struct VerifiedArtifacts {
131    pub model_dir: PathBuf,
132    pub backend_kind: NerBackendKind,
133    pub recognizer_model_id: String,
134    pub recognizer_model_version: String,
135    pub labels: LabelMap,
136    pub id2label: Vec<String>,
137}
gaze_recognizers/ner/types.rs

gaze_recognizers/ner/
types.rs
