gaze_recognizers/safety_net/kiji_distilbert/
mod.rs

1//! Kiji DistilBERT safety-net adapter (v0.8 Tier 2.5).
2//!
3//! Pass-3 SafetyNet backend that runs a pinned DistilBERT NER model as a
4//! subprocess. The contract is observer-only: every label arrives as a
5//! validated `KijiLabel`, every span passes through `normalize_raw_spans`,
6//! and no PII-bearing upstream field crosses the adapter boundary. This is
7//! the second NER opinion at the chokepoint (Axis 1 — defense in depth).
8//!
9//! The subprocess reads tokenized clean text on stdin and emits a JSON span
10//! array on stdout (`[{label, start, end, score?}, ...]`). The `_text` and
11//! `_placeholder` fields, if present, are dropped at deserialization with a
12//! buffer-scrub on `Drop`.
13
14use std::sync::{Arc, OnceLock};
15
16use gaze_types::{LeakSuspect, LocaleTag, SafetyNet, SafetyNetContext, SafetyNetError};
17
18use crate::{LocaleAwareModel, ModelError, ModelHints, ModelInput, ModelSpan};
19
20pub mod backend;
21pub mod class_map;
22
23pub use backend::artifacts::{
24    KIJI_DISTILBERT_BUNDLE_SHA256, KIJI_DISTILBERT_HF_COMMIT, KIJI_DISTILBERT_HF_REPO,
25    KIJI_DISTILBERT_INT8_BUNDLE_SHA256, KIJI_DISTILBERT_INT8_SHA256SUMS,
26    KIJI_DISTILBERT_SHA256SUMS, REQUIRED_KIJI_ARTIFACTS, REQUIRED_KIJI_INT8_ARTIFACTS,
27};
28#[cfg(feature = "runtime-candle")]
29pub use backend::candle::{CandleKijiBackend, CandleKijiConfig};
30pub use backend::ort::{OrtKijiBackend, OrtKijiConfig};
31pub use backend::subprocess::{SubprocessKijiBackend, SubprocessKijiConfig};
32#[cfg(feature = "runtime-tract")]
33pub use backend::tract::{TractKijiBackend, TractKijiConfig};
34pub use backend::{KijiBackendKind, KijiDistilbertBackend, KijiDistilbertPrecision, RawSpan};
35pub use class_map::{
36    kiji_label_to_pii_class, kiji_label_to_safety_net_class, map_kiji_label, validate_kiji_label,
37    KijiLabel,
38};
39
40pub type KijiOrtBackend = OrtKijiBackend;
41pub type KijiOrtConfig = OrtKijiConfig;
42
43/// Safety-net implementation backed by a lazily initialized Kiji subprocess.
44///
45/// Same caching strategy as `OpenAiFilterSafetyNet`: deterministic init
46/// failures (missing SHA256SUMS, missing artifacts, bad perms) cache in a
47/// `OnceLock` so we never retry-storm on every check.
48pub struct KijiDistilbertSafetyNet {
49    locales: Vec<LocaleTag>,
50    config: KijiDistilbertConfig,
51    backend: OnceLock<Result<Arc<dyn KijiDistilbertBackend>, Arc<SafetyNetError>>>,
52}
53
54impl std::fmt::Debug for KijiDistilbertSafetyNet {
55    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56        formatter
57            .debug_struct("KijiDistilbertSafetyNet")
58            .field("locales", &self.locales)
59            .field("backend_kind", &self.backend_kind())
60            .finish_non_exhaustive()
61    }
62}
63
64impl KijiDistilbertSafetyNet {
65    pub fn new(config: impl Into<KijiDistilbertConfig>) -> Self {
66        Self {
67            locales: vec![LocaleTag::Global],
68            config: config.into(),
69            backend: OnceLock::new(),
70        }
71    }
72
73    pub fn from_env() -> Result<Self, SafetyNetError> {
74        Ok(Self::new(SubprocessKijiConfig::from_env()?))
75    }
76
77    pub fn from_env_ort() -> Result<Self, SafetyNetError> {
78        Ok(Self::new(OrtKijiConfig::from_env()?))
79    }
80
81    pub fn backend_kind(&self) -> KijiBackendKind {
82        self.config.kind()
83    }
84
85    pub fn with_locales(mut self, locales: Vec<LocaleTag>) -> Self {
86        self.locales = locales;
87        self
88    }
89
90    fn backend(&self) -> Result<Arc<dyn KijiDistilbertBackend>, SafetyNetError> {
91        match self.backend.get_or_init(|| self.config.load_backend()) {
92            Ok(backend) => Ok(Arc::clone(backend)),
93            Err(error) => Err((**error).clone()),
94        }
95    }
96}
97
98#[derive(Debug, Clone)]
99pub enum KijiDistilbertConfig {
100    Subprocess(SubprocessKijiConfig),
101    Ort(OrtKijiConfig),
102    #[cfg(feature = "runtime-tract")]
103    Tract(TractKijiConfig),
104    #[cfg(feature = "runtime-candle")]
105    Candle(CandleKijiConfig),
106}
107
108impl KijiDistilbertConfig {
109    pub fn kind(&self) -> KijiBackendKind {
110        match self {
111            Self::Subprocess(_) => KijiBackendKind::Subprocess,
112            Self::Ort(_) => KijiBackendKind::Ort,
113            #[cfg(feature = "runtime-tract")]
114            Self::Tract(_) => KijiBackendKind::Tract,
115            #[cfg(feature = "runtime-candle")]
116            Self::Candle(_) => KijiBackendKind::Candle,
117        }
118    }
119
120    fn load_backend(&self) -> Result<Arc<dyn KijiDistilbertBackend>, Arc<SafetyNetError>> {
121        match self {
122            Self::Subprocess(config) => SubprocessKijiBackend::new(config.clone())
123                .map(|backend| Arc::new(backend) as Arc<dyn KijiDistilbertBackend>)
124                .map_err(Arc::new),
125            Self::Ort(config) => OrtKijiBackend::new(config.clone())
126                .map(|backend| Arc::new(backend) as Arc<dyn KijiDistilbertBackend>)
127                .map_err(Arc::new),
128            #[cfg(feature = "runtime-tract")]
129            Self::Tract(config) => TractKijiBackend::new(config.clone())
130                .map(|backend| Arc::new(backend) as Arc<dyn KijiDistilbertBackend>)
131                .map_err(Arc::new),
132            #[cfg(feature = "runtime-candle")]
133            Self::Candle(config) => CandleKijiBackend::new(config.clone())
134                .map(|backend| Arc::new(backend) as Arc<dyn KijiDistilbertBackend>)
135                .map_err(Arc::new),
136        }
137    }
138}
139
140impl From<SubprocessKijiConfig> for KijiDistilbertConfig {
141    fn from(config: SubprocessKijiConfig) -> Self {
142        Self::Subprocess(config)
143    }
144}
145
146impl From<OrtKijiConfig> for KijiDistilbertConfig {
147    fn from(config: OrtKijiConfig) -> Self {
148        Self::Ort(config)
149    }
150}
151
152#[cfg(feature = "runtime-tract")]
153impl From<TractKijiConfig> for KijiDistilbertConfig {
154    fn from(config: TractKijiConfig) -> Self {
155        Self::Tract(config)
156    }
157}
158
159#[cfg(feature = "runtime-candle")]
160impl From<CandleKijiConfig> for KijiDistilbertConfig {
161    fn from(config: CandleKijiConfig) -> Self {
162        Self::Candle(config)
163    }
164}
165
166impl SafetyNet for KijiDistilbertSafetyNet {
167    fn id(&self) -> &str {
168        "kiji-distilbert"
169    }
170
171    fn supported_locales(&self) -> &[LocaleTag] {
172        &self.locales
173    }
174
175    fn check(
176        &self,
177        clean_text: &str,
178        context: SafetyNetContext<'_>,
179    ) -> Result<Vec<LeakSuspect>, SafetyNetError> {
180        let backend = self.backend()?;
181        let raw_spans = backend.infer(clean_text)?;
182        raw_spans
183            .into_iter()
184            .filter_map(|raw| match raw_span_to_suspect(raw, &*backend, context) {
185                Ok(Some(suspect)) => Some(Ok(suspect)),
186                Ok(None) => None,
187                Err(error) => Some(Err(error)),
188            })
189            .collect()
190    }
191}
192
193impl LocaleAwareModel for KijiDistilbertSafetyNet {
194    fn name(&self) -> &str {
195        "kiji-distilbert"
196    }
197
198    fn native_locales(&self) -> &[LocaleTag] {
199        &self.locales
200    }
201
202    fn infer(&self, input: ModelInput, hints: ModelHints) -> Result<Vec<ModelSpan>, ModelError> {
203        let backend = self
204            .backend()
205            .map_err(|error| ModelError::InitFailed(error.to_string()))?;
206        let mut spans = backend
207            .infer(&input.text)
208            .map_err(|error| ModelError::InferenceFailed(error.to_string()))?
209            .into_iter()
210            .map(|raw| raw_span_to_model_span(raw, self.name(), &input.text))
211            .collect::<Result<Vec<_>, _>>()?;
212
213        if let Some(max_spans) = hints.max_spans {
214            spans.truncate(max_spans as usize);
215        }
216
217        Ok(spans)
218    }
219}
220
221fn raw_span_to_model_span(
222    raw: RawSpan,
223    model_name: &str,
224    input: &str,
225) -> Result<ModelSpan, ModelError> {
226    if raw.start >= raw.end
227        || raw.end > input.len()
228        || !input.is_char_boundary(raw.start)
229        || !input.is_char_boundary(raw.end)
230    {
231        return Err(ModelError::InferenceFailed(
232            "kiji returned out-of-bounds span".to_string(),
233        ));
234    }
235
236    let label = map_kiji_label(&raw.label)
237        .map_err(|error| ModelError::InferenceFailed(error.to_string()))?;
238    let class = kiji_label_to_pii_class(label);
239    let byte_range = raw.start..raw.end;
240    let text = input[byte_range.clone()].to_string();
241
242    Ok(ModelSpan {
243        text,
244        byte_range,
245        class,
246        confidence: raw.score,
247        model_name: model_name.to_string(),
248    })
249}
250
251fn raw_span_to_suspect(
252    raw: RawSpan,
253    backend: &dyn KijiDistilbertBackend,
254    context: SafetyNetContext<'_>,
255) -> Result<Option<LeakSuspect>, SafetyNetError> {
256    let label = map_kiji_label(&raw.label)?;
257    let class = kiji_label_to_pii_class(label);
258    let span = raw.start..raw.end;
259    let Some(kind) = context.manifest.diff_against(&span, &class) else {
260        return Ok(None);
261    };
262
263    Ok(Some(LeakSuspect::new(
264        span,
265        class,
266        backend.id(),
267        raw.score,
268        kind,
269        raw.label,
270        context.field_path.map(str::to_string),
271    )))
272}
273
274#[cfg(test)]
275mod tests {
276    use gaze_types::{DocumentKind, Manifest};
277
278    use super::*;
279
280    #[test]
281    fn empty_command_failure_is_cached() {
282        let net = KijiDistilbertSafetyNet::new(SubprocessKijiConfig::new(""));
283        let manifest = Manifest::default();
284        let context = SafetyNetContext::new(
285            &manifest,
286            &[LocaleTag::Global],
287            DocumentKind::Text,
288            None,
289            None,
290        );
291
292        let first = net.check("clean", context).unwrap_err();
293        let second = net.check("clean", context).unwrap_err();
294        assert_eq!(first, second);
295        assert!(net.backend.get().is_some());
296    }
297}
gaze_recognizers/safety_net/kiji_distilbert/mod.rs

gaze_recognizers/safety_net/kiji_distilbert/
mod.rs
