| use anyhow::{Context, Result, bail, ensure}; |
| use ort::{session::Session, value::Tensor}; |
| use serde::{Deserialize, Serialize}; |
| use serde_json::json; |
| use std::{ |
| io::{self, BufRead, Write}, |
| path::Path, |
| }; |
| use tokenizers::{Encoding, Tokenizer}; |
|
|
| #[derive(Deserialize)] |
| struct Config { |
| labels: Vec<String>, |
| max_length: usize, |
| stride: usize, |
| #[serde(default)] |
| calibration: Option<Calibration>, |
| #[serde(default)] |
| calibration_file: Option<String>, |
| } |
| #[derive(Deserialize)] |
| struct Calibration { |
| #[serde(alias = "xs")] |
| x: Vec<f64>, |
| #[serde(alias = "ys")] |
| y: Vec<f64>, |
| threshold: Option<f64>, |
| } |
| impl Calibration { |
| fn apply(&self, v: f64) -> f64 { |
| let i = self.x.partition_point(|x| *x < v); |
| if i == 0 { |
| self.y[0] |
| } else if i == self.x.len() { |
| *self.y.last().unwrap() |
| } else { |
| self.y[i - 1] |
| + (v - self.x[i - 1]) / (self.x[i] - self.x[i - 1]) * (self.y[i] - self.y[i - 1]) |
| } |
| } |
| } |
| #[derive(Clone, Debug, Serialize)] |
| struct Finding { |
| label: String, |
| start: usize, |
| end: usize, |
| raw_score: f64, |
| confidence: Option<f64>, |
| #[serde(skip)] |
| boundary_fragment: bool, |
| } |
| fn encode_windows( |
| tokenizer: &Tokenizer, |
| text: &str, |
| max_length: usize, |
| stride: usize, |
| ) -> Result<Vec<Encoding>> { |
| let full = tokenizer |
| .encode(text, false) |
| .map_err(|e| anyhow::anyhow!("{e}"))?; |
| let capacity = max_length - 2; |
| ensure!(stride < capacity, "invalid stride"); |
| let cls = tokenizer |
| .token_to_id("[CLS]") |
| .context("missing CLS token")?; |
| let sep = tokenizer |
| .token_to_id("[SEP]") |
| .context("missing SEP token")?; |
| let mut windows = Vec::new(); |
| let mut start = 0; |
| loop { |
| let end = (start + capacity).min(full.len()); |
| let n = end - start + 2; |
| let mut ids = vec![cls]; |
| ids.extend_from_slice(&full.get_ids()[start..end]); |
| ids.push(sep); |
| let mut offsets = vec![(0, 0)]; |
| offsets.extend_from_slice(&full.get_offsets()[start..end]); |
| offsets.push((0, 0)); |
| let mut tokens = vec!["[CLS]".to_string()]; |
| tokens.extend_from_slice(&full.get_tokens()[start..end]); |
| tokens.push("[SEP]".to_string()); |
| let mut special = vec![0; n]; |
| special[0] = 1; |
| special[n - 1] = 1; |
| windows.push(Encoding::new( |
| ids, |
| vec![0; n], |
| tokens, |
| vec![None; n], |
| offsets, |
| special, |
| vec![1; n], |
| vec![], |
| Default::default(), |
| )); |
| if end == full.len() { |
| break; |
| } |
| start += capacity - stride; |
| } |
| Ok(windows) |
| } |
| fn decode(enc: &tokenizers::Encoding, logits: &[f32], labels: &[String]) -> Vec<Finding> { |
| let k = 1 + 2 * labels.len(); |
| let mut out = Vec::new(); |
| let mut current: Option<Finding> = None; |
| for (i, row) in logits.chunks_exact(k).enumerate() { |
| let (start, end) = enc.get_offsets()[i]; |
| if enc.get_special_tokens_mask()[i] != 0 || enc.get_attention_mask()[i] == 0 || start == end |
| { |
| if let Some(f) = current.take() { |
| out.push(f); |
| } |
| continue; |
| } |
| let mut id = 0; |
| for j in 1..k { |
| if row[j] > row[id] { |
| id = j; |
| } |
| } |
| if id == 0 { |
| if let Some(f) = current.take() { |
| out.push(f); |
| } |
| continue; |
| } |
| let label = &labels[(id - 1) / 2]; |
| let max = row[id] as f64; |
| let prob = 1.0 / row.iter().map(|v| (*v as f64 - max).exp()).sum::<f64>(); |
| if id % 2 == 0 && current.as_ref().is_some_and(|f| f.label == *label) { |
| let f = current.as_mut().unwrap(); |
| f.end = end; |
| f.raw_score = f.raw_score.min(prob); |
| } else { |
| if let Some(f) = current.take() { |
| out.push(f); |
| } |
| current = Some(Finding { |
| label: label.clone(), |
| start, |
| end, |
| raw_score: prob, |
| confidence: None, |
| boundary_fragment: false, |
| }); |
| } |
| } |
| if let Some(f) = current { |
| out.push(f); |
| } |
| out |
| } |
| fn mark_boundaries( |
| enc: &tokenizers::Encoding, |
| findings: &mut [Finding], |
| has_previous: bool, |
| has_next: bool, |
| ) { |
| let content: Vec<_> = enc.get_offsets().iter().filter(|(a, b)| a != b).collect(); |
| if let (Some(first), Some(last)) = (content.first(), content.last()) { |
| for f in findings { |
| f.boundary_fragment = |
| (has_previous && f.start == first.0) || (has_next && f.end == last.1); |
| } |
| } |
| } |
| fn resolve(findings: Vec<Finding>) -> Vec<Finding> { |
| let mut findings: Vec<Finding> = findings |
| .iter() |
| .filter(|f| { |
| !(f.boundary_fragment |
| && findings.iter().any(|g| { |
| g.label == f.label |
| && g.start <= f.start |
| && g.end >= f.end |
| && (g.start < f.start || g.end > f.end) |
| })) |
| }) |
| .cloned() |
| .collect(); |
| findings.sort_by(|a, b| { |
| b.raw_score |
| .total_cmp(&a.raw_score) |
| .then(a.start.cmp(&b.start)) |
| .then(a.end.cmp(&b.end)) |
| .then(a.label.cmp(&b.label)) |
| }); |
| let mut accepted: Vec<Finding> = Vec::new(); |
| for f in findings { |
| if !accepted.iter().any(|a| f.start < a.end && a.start < f.end) { |
| accepted.push(f); |
| } |
| } |
| accepted.sort_by_key(|f| (f.start, f.end)); |
| accepted |
| } |
| fn redact(text: &str, findings: &[Finding]) -> Result<String> { |
| let mut s = String::new(); |
| let mut pos = 0; |
| for f in findings { |
| ensure!( |
| f.start >= pos |
| && f.end <= text.len() |
| && text.is_char_boundary(f.start) |
| && text.is_char_boundary(f.end), |
| "invalid UTF-8 finding boundary" |
| ); |
| s.push_str(&text[pos..f.start]); |
| s.push_str("[REDACTED]"); |
| pos = f.end; |
| } |
| s.push_str(&text[pos..]); |
| Ok(s) |
| } |
| struct Detector { |
| config: Config, |
| tokenizer: Tokenizer, |
| session: Session, |
| } |
| impl Detector { |
| fn load(path: &Path) -> Result<Self> { |
| let mut config: Config = serde_json::from_slice(&std::fs::read(path.join("config.json"))?)?; |
| if let Some(file) = &config.calibration_file { |
| config.calibration = Some(serde_json::from_slice(&std::fs::read(path.join(file))?)?); |
| } |
| ensure!(config.labels.len() == 7, "expected seven labels"); |
| ensure!( |
| config.max_length > 2 && config.stride < config.max_length - 2, |
| "invalid chunk settings" |
| ); |
| if let Some(c) = &config.calibration { |
| ensure!( |
| !c.x.is_empty() && c.x.len() == c.y.len(), |
| "invalid calibration lengths" |
| ); |
| ensure!( |
| c.x.iter().all(|x| x.is_finite()) |
| && c.y.iter().all(|y| y.is_finite() && (0.0..=1.0).contains(y)) |
| && c.x.windows(2).all(|v| v[0] < v[1]) |
| && c.y.windows(2).all(|v| v[0] <= v[1]), |
| "invalid calibration values" |
| ); |
| } |
| let mut tokenizer = Tokenizer::from_file(path.join("tokenizer.json")) |
| .map_err(|e| anyhow::anyhow!("{e}"))?; |
| tokenizer.with_padding(None); |
| tokenizer |
| .with_truncation(None) |
| .map_err(|e| anyhow::anyhow!("{e}"))?; |
| let _ = ort::init_from(path.join(if cfg!(target_os = "windows") { "onnxruntime.dll" } else if cfg!(target_os = "macos") { "libonnxruntime.dylib" } else { "libonnxruntime.so" }))?.commit(); |
| let session = Session::builder()? |
| .with_intra_threads(1) |
| .map_err(|e| anyhow::anyhow!("{e}"))? |
| .commit_from_file(path.join("model.onnx"))?; |
| Ok(Self { |
| config, |
| tokenizer, |
| session, |
| }) |
| } |
| fn predict(&mut self, text: &str, debug: bool) -> Result<serde_json::Value> { |
| let encodings = encode_windows( |
| &self.tokenizer, |
| text, |
| self.config.max_length, |
| self.config.stride, |
| )?; |
| let mut findings = Vec::new(); |
| let mut windows = Vec::new(); |
| for (index, e) in encodings.iter().enumerate() { |
| let n = e.len(); |
| let ids: Vec<i64> = e.get_ids().iter().map(|x| *x as i64).collect(); |
| let mask: Vec<i64> = e.get_attention_mask().iter().map(|x| *x as i64).collect(); |
| let outputs=self.session.run(ort::inputs!["input_ids"=>Tensor::from_array(([1,n],ids.clone()))?,"attention_mask"=>Tensor::from_array(([1,n],mask.clone()))?])?; |
| let (shape, logits) = outputs["logits"].try_extract_tensor::<f32>()?; |
| ensure!( |
| shape.as_ref() == [1, n as i64, 15], |
| "unexpected logits shape: {shape:?}" |
| ); |
| ensure!(logits.iter().all(|x| x.is_finite()), "non-finite logits"); |
| let mut decoded = decode(e, logits, &self.config.labels); |
| mark_boundaries(e, &mut decoded, index > 0, index + 1 < encodings.len()); |
| if n == self.config.max_length { |
| let content: Vec<_> = e.get_offsets().iter().filter(|(a, b)| a != b).collect(); |
| if let (Some(first), Some(last)) = (content.first(), content.last()) { |
| ensure!( |
| !decoded |
| .iter() |
| .any(|f| f.start == first.0 && f.end == last.1), |
| "entity spans a full window; complete entity boundary is unsupported" |
| ); |
| } |
| } |
| findings.extend(decoded); |
| if debug { |
| windows.push(json!({"input_ids":ids,"attention_mask":mask,"offsets":e.get_offsets(),"special_tokens_mask":e.get_special_tokens_mask(),"logits":logits.chunks_exact(15).collect::<Vec<_>>()})); |
| } |
| } |
| let mut findings = resolve(findings); |
| if let Some(c) = &self.config.calibration { |
| for f in &mut findings { |
| f.confidence = Some(c.apply(f.raw_score)); |
| } |
| if let Some(t) = c.threshold { |
| findings.retain(|f| f.confidence.unwrap() >= t); |
| } |
| } |
| let redacted = redact(text, &findings)?; |
| let mut result = |
| json!({"findings":findings,"contains_pii":!findings.is_empty(),"redacted":redacted}); |
| if debug { |
| result["windows"] = json!(windows); |
| } |
| Ok(result) |
| } |
| } |
| fn main() -> Result<()> { |
| let args: Vec<String> = std::env::args().skip(1).collect(); |
| if args.is_empty() { |
| bail!("usage: datafog-pii BUNDLE [--debug]; stdin JSONL objects with text:string"); |
| } |
| ensure!( |
| args.len() == 1 || (args.len() == 2 && args[1] == "--debug"), |
| "unknown arguments" |
| ); |
| let mut detector = Detector::load(Path::new(&args[0])).context("loading detector")?; |
| let debug = args.len() == 2; |
| let stdin = io::stdin(); |
| let mut stdout = io::BufWriter::new(io::stdout().lock()); |
| for line in stdin.lock().lines() { |
| let line = line?; |
| let result = (|| -> Result<_> { |
| let input: serde_json::Value = serde_json::from_str(&line)?; |
| let text = input |
| .get("text") |
| .and_then(|t| t.as_str()) |
| .context("expected text:string")?; |
| detector.predict(text, debug) |
| })(); |
| let output = match result { |
| Ok(v) => v, |
| Err(e) => json!({"error":format!("{e:#}")}), |
| }; |
| writeln!(stdout, "{}", serde_json::to_string(&output)?)?; |
| stdout.flush()?; |
| } |
| Ok(()) |
| } |
| #[cfg(test)] |
| mod tests { |
| use super::*; |
| #[test] |
| fn unicode_overlap_and_adjacent() { |
| let f = |start, end, score| Finding { |
| label: "first_name".into(), |
| start, |
| end, |
| raw_score: score, |
| confidence: None, |
| boundary_fragment: false, |
| }; |
| let spans = resolve(vec![f(5, 8, 0.6), f(5, 6, 0.9), f(6, 8, 0.8), f(5, 6, 0.7)]); |
| assert_eq!(spans.len(), 2); |
| assert_eq!( |
| redact("🙂 José", &spans).unwrap(), |
| "🙂 [REDACTED][REDACTED]é" |
| ); |
| assert!(redact("🙂", &[f(1, 3, 1.)]).is_err()); |
| } |
| #[test] |
| fn oracle_boundaries_and_complete_token_coverage() { |
| let fixture: serde_json::Value = |
| serde_json::from_str(include_str!("../tests/boundary-oracle.json")).unwrap(); |
| let mut tokenizer = Tokenizer::from_file( |
| Path::new(env!("CARGO_MANIFEST_DIR")).join("../bundle/tokenizer.json"), |
| ) |
| .unwrap(); |
| tokenizer.with_truncation(None).unwrap(); |
| tokenizer.with_padding(None); |
| let labels: Vec<String> = serde_json::from_value(fixture["labels"].clone()).unwrap(); |
| for case in fixture["cases"].as_array().unwrap() { |
| let text = case["text"].as_str().unwrap(); |
| let windows = encode_windows(&tokenizer, text, 512, 42).unwrap(); |
| assert_eq!(windows.len(), case["windows"].as_array().unwrap().len()); |
| let mut findings = Vec::new(); |
| let mut covered = Vec::new(); |
| for (i, e) in windows.iter().enumerate() { |
| let oracle = &case["windows"][i]; |
| assert_eq!( |
| serde_json::to_value(e.get_ids()).unwrap(), |
| oracle["input_ids"] |
| ); |
| assert_eq!( |
| serde_json::to_value(e.get_offsets()).unwrap(), |
| oracle["offsets"] |
| ); |
| covered.extend_from_slice(&e.get_ids()[if i == 0 { 1 } else { 43 }..e.len() - 1]); |
| let logits: Vec<Vec<f32>> = |
| serde_json::from_value(oracle["logits"].clone()).unwrap(); |
| let mut f = decode(e, &logits.concat(), &labels); |
| mark_boundaries(e, &mut f, i > 0, i + 1 < windows.len()); |
| findings.extend(f); |
| } |
| let full = tokenizer.encode(text, false).unwrap(); |
| assert_eq!(covered, full.get_ids()); |
| let actual = resolve(findings); |
| let spans: Vec<_> = actual |
| .iter() |
| .map(|f| json!({"label":f.label,"start":f.start,"end":f.end})) |
| .collect(); |
| assert_eq!(json!(spans), case["expected"]); |
| assert_eq!( |
| redact(text, &actual).unwrap(), |
| case["redacted"].as_str().unwrap() |
| ); |
| } |
| } |
| #[test] |
| fn interpolation() { |
| let c = Calibration { |
| x: vec![0.2, 0.8], |
| y: vec![0.1, 0.9], |
| threshold: None, |
| }; |
| assert!((c.apply(0.5) - 0.5).abs() < 1e-12); |
| assert_eq!(c.apply(0.), 0.1); |
| assert_eq!(c.apply(1.), 0.9); |
| } |
| } |
|
|
