Skip to main content

infer/
infer.rs

1//! Example: run privacy filter inference on sample text.
2//!
3//! ```bash
4//! cargo run --example infer --release -- --model-dir /path/to/privacy-filter
5//! ```
6
7use std::path::PathBuf;
8use clap::Parser;
9use privacy_filter_rs::backend::{B, Device};
10
11#[derive(Parser)]
12struct Args {
13    #[arg(short = 'm', long)]
14    model_dir: PathBuf,
15
16    #[arg(short = 't', long, default_value = "0")]
17    threads: usize,
18}
19
20fn main() -> anyhow::Result<()> {
21    let args = Args::parse();
22
23    let n = privacy_filter_rs::init_threads(Some(args.threads));
24    eprintln!("Using {n} threads");
25
26    let device = <Device as Default>::default();
27
28    eprintln!("Loading model...");
29    let engine = privacy_filter_rs::PrivacyFilterInference::<B>::load(
30        &args.model_dir,
31        device,
32    )?;
33
34    let samples = [
35        "My name is Alice Smith and I live at 123 Main Street, Springfield.",
36        "You can reach me at alice.smith@example.com or call 555-0123.",
37        "My account number is 4532-1234-5678-9012 and my password is hunter2.",
38        "Born on January 15, 1990, Alice visited https://secret-site.com/login.",
39        "The weather is nice today and the stock market went up.",
40    ];
41
42    for text in &samples {
43        println!("\n--- Input: {text}");
44        let spans = engine.predict(text)?;
45        if spans.is_empty() {
46            println!("  No PII detected.");
47        } else {
48            for span in &spans {
49                println!(
50                    "  [{:>15}] {:<30} (score: {:.4}, chars {}..{})",
51                    span.entity_group, span.word, span.score, span.start, span.end
52                );
53            }
54        }
55    }
56
57    Ok(())
58}