privacy-filter-rs 0.1.0

OpenAI Privacy Filter — PII detection inference in pure Rust with Burn ML
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87

//! # privacy-filter-rs — OpenAI Privacy Filter inference in Rust
//!
//! Pure-Rust inference for the [OpenAI Privacy Filter](https://huggingface.co/openai/privacy-filter)
//! token classification model, built on [Burn 0.20](https://burn.dev).
//!
//! ## Quick start
//!
//! ```rust,ignore
//! use privacy_filter_rs::PrivacyFilterInference;
//! use std::path::Path;
//!
//! // Choose backend
//! type B = burn::backend::NdArray;
//! let device = burn::backend::ndarray::NdArrayDevice::Cpu;
//!
//! let engine = PrivacyFilterInference::<B>::load(
//!     Path::new("/path/to/privacy-filter"),
//!     device,
//! )?;
//!
//! let spans = engine.predict("My name is Alice Smith")?;
//! for span in &spans {
//!     println!("{}: {} (score: {:.4})", span.entity_group, span.word, span.score);
//! }
//! ```

// ── Thread configuration ────────────────────────────────────────────────────

/// Configure the global Rayon thread pool.
///
/// Call this **once**, before any model operations.
/// Returns the actual number of threads in the pool.
pub fn init_threads(n: Option<usize>) -> usize {
    let mut builder = rayon::ThreadPoolBuilder::new();
    if let Some(count) = n {
        if count > 0 {
            builder = builder.num_threads(count);
        }
    }
    let _ = builder.build_global();
    rayon::current_num_threads()
}

// ── Internal modules ────────────────────────────────────────────────────────

pub mod config;
pub mod inference;
pub mod model;
pub mod tensor_utils;
pub mod viterbi;
pub mod weights;

// ── Flat re-exports ─────────────────────────────────────────────────────────

pub use config::{ModelConfig, ViterbiConfig};
pub use inference::PrivacyFilterInference;
pub use viterbi::PrivacySpan;

// ── Backend selection ───────────────────────────────────────────────────────

#[cfg(feature = "ndarray")]
pub mod backend {
    pub use burn::backend::NdArray as B;
    pub type Device = burn::backend::ndarray::NdArrayDevice;
    pub fn device() -> Device { Device::Cpu }
}

#[cfg(all(feature = "wgpu-f16", not(feature = "ndarray"), not(feature = "wgpu"), not(feature = "mlx")))]
pub mod backend {
    pub type B = burn::backend::wgpu::Wgpu<half::f16, i32, u32>;
    pub type Device = burn::backend::wgpu::WgpuDevice;
    pub fn device() -> Device { Device::DefaultDevice }
}

#[cfg(all(feature = "wgpu", not(feature = "ndarray"), not(feature = "wgpu-f16"), not(feature = "mlx")))]
pub mod backend {
    pub use burn::backend::Wgpu as B;
    pub type Device = burn::backend::wgpu::WgpuDevice;
    pub fn device() -> Device { Device::DefaultDevice }
}

#[cfg(all(feature = "mlx", not(feature = "ndarray"), not(feature = "wgpu"), not(feature = "wgpu-f16")))]
pub mod backend {
    pub use burn_mlx::Mlx as B;
    pub type Device = burn_mlx::MlxDevice;
    pub fn device() -> Device { Default::default() }
}