gaze_recognizers/
locale_aware.rs1use std::ops::Range;
2
3use gaze_types::{LocaleTag, PiiClass};
4use thiserror::Error;
5
6pub trait LocaleAwareModel: Send + Sync {
8 fn name(&self) -> &str;
10
11 fn native_locales(&self) -> &[LocaleTag];
13
14 fn infer(&self, input: ModelInput, hints: ModelHints) -> Result<Vec<ModelSpan>, ModelError>;
16}
17
18#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct ModelInput {
21 pub text: String,
22 pub locale: LocaleTag,
23}
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub struct ModelHints {
28 pub stage: ModelStage,
29 pub max_spans: Option<u32>,
30}
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum ModelStage {
35 Pass2Ner,
36 Pass3SafetyNet,
37}
38
39#[derive(Debug, Clone, PartialEq)]
41pub struct ModelSpan {
42 pub text: String,
43 pub byte_range: Range<usize>,
44 pub class: PiiClass,
45 pub confidence: Option<f32>,
46 pub model_name: String,
47}
48
49#[derive(Debug, Clone, PartialEq, Eq, Error)]
51pub enum ModelError {
52 #[error("backend init failed: {0}")]
53 InitFailed(String),
54 #[error("inference failed: {0}")]
55 InferenceFailed(String),
56 #[error("locale not supported: {0:?}")]
57 LocaleNotSupported(LocaleTag),
58 #[error("model integrity mismatch")]
59 IntegrityMismatch,
60 #[error("no locale model coverage for {locale:?} at {stage:?}")]
61 NoLocaleModelCoverage {
62 locale: LocaleTag,
63 stage: ModelStage,
64 },
65 #[error("backend internal error: {0}")]
66 Internal(String),
67}
68
69#[derive(Default)]
71pub struct LocaleAwareModelRegistry {
72 backends: Vec<Box<dyn LocaleAwareModel>>,
73}
74
75impl LocaleAwareModelRegistry {
76 pub fn new() -> Self {
77 Self::default()
78 }
79
80 pub fn from_backends(backends: Vec<Box<dyn LocaleAwareModel>>) -> Self {
81 Self { backends }
82 }
83
84 pub fn register(&mut self, backend: impl LocaleAwareModel + 'static) {
85 self.backends.push(Box::new(backend));
86 }
87
88 pub fn len(&self) -> usize {
89 self.backends.len()
90 }
91
92 pub fn is_empty(&self) -> bool {
93 self.backends.is_empty()
94 }
95
96 pub fn resolve(
99 &self,
100 locale: &LocaleTag,
101 stage: ModelStage,
102 ) -> Result<Vec<&dyn LocaleAwareModel>, ModelError> {
103 if let Some(matches) = self.matches_for(locale) {
104 return Ok(matches);
105 }
106
107 if let Some(parent) = parent_language(locale) {
108 if let Some(matches) = self.matches_for(&parent) {
109 return Ok(matches);
110 }
111 }
112
113 if let Some(matches) = self.matches_for(&LocaleTag::Global) {
114 return Ok(matches);
115 }
116
117 Err(ModelError::NoLocaleModelCoverage {
118 locale: locale.clone(),
119 stage,
120 })
121 }
122
123 fn matches_for(&self, locale: &LocaleTag) -> Option<Vec<&dyn LocaleAwareModel>> {
124 let matches: Vec<&dyn LocaleAwareModel> = self
125 .backends
126 .iter()
127 .filter(|backend| backend.native_locales().contains(locale))
128 .map(|backend| backend.as_ref())
129 .collect();
130
131 (!matches.is_empty()).then_some(matches)
132 }
133}
134
135fn parent_language(locale: &LocaleTag) -> Option<LocaleTag> {
136 match locale {
137 LocaleTag::DeDe | LocaleTag::DeAt | LocaleTag::DeCh => {
138 Some(LocaleTag::Other("de".to_string()))
139 }
140 LocaleTag::EnUs | LocaleTag::EnGb | LocaleTag::EnIe | LocaleTag::EnAu | LocaleTag::EnCa => {
141 Some(LocaleTag::Other("en".to_string()))
142 }
143 LocaleTag::Other(raw) => raw
144 .split_once('-')
145 .map(|(language, _)| LocaleTag::Other(language.to_ascii_lowercase())),
146 LocaleTag::Global => None,
147 _ => None,
148 }
149}