use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use gaze_types::SafetyNetError;
use tract_onnx::prelude::*;
use super::artifacts::{verify_model_dir, KIJI_DISTILBERT_BUNDLE_SHA256};
use super::decode::decode_logits;
use super::{normalize_raw_spans, KijiDistilbertBackend, RawSpan};
const DEFAULT_MAX_INPUT_BYTES: usize = 1024 * 1024;
const MODEL_FILE: &str = "model.onnx";
const TOKENIZER_FILE: &str = "tokenizer.json";
#[derive(Debug, Clone)]
pub struct TractKijiConfig {
model_dir: PathBuf,
max_input_bytes: usize,
version: String,
decoding_params: Vec<(&'static str, String)>,
#[cfg(any(test, feature = "test-support"))]
expected_bundle_sha256: String,
}
impl TractKijiConfig {
pub fn new(model_dir: impl Into<PathBuf>) -> Self {
Self {
model_dir: model_dir.into(),
max_input_bytes: DEFAULT_MAX_INPUT_BYTES,
version: "kiji/distilbert:tract".to_string(),
decoding_params: vec![("runtime", "tract".to_string())],
#[cfg(any(test, feature = "test-support"))]
expected_bundle_sha256: KIJI_DISTILBERT_BUNDLE_SHA256.to_string(),
}
}
pub fn with_max_input_bytes(mut self, max_input_bytes: usize) -> Self {
self.max_input_bytes = max_input_bytes;
self
}
pub fn model_dir(&self) -> &Path {
&self.model_dir
}
fn expected_bundle_sha256(&self) -> &str {
#[cfg(any(test, feature = "test-support"))]
{
&self.expected_bundle_sha256
}
#[cfg(not(any(test, feature = "test-support")))]
{
KIJI_DISTILBERT_BUNDLE_SHA256
}
}
}
#[derive(Debug)]
pub struct TractKijiBackend {
config: TractKijiConfig,
tokenizer: tokenizers::Tokenizer,
model: Mutex<Arc<TypedRunnableModel>>,
has_token_type_ids: bool,
}
impl TractKijiBackend {
pub fn new(config: TractKijiConfig) -> Result<Self, SafetyNetError> {
verify_model_dir(Some(config.model_dir()), config.expected_bundle_sha256())?;
let tokenizer = tokenizers::Tokenizer::from_file(config.model_dir.join(TOKENIZER_FILE))
.map_err(|err| SafetyNetError::ModelUnavailable {
reason: format!(
"failed to load kiji tokenizer: {}",
sanitize_error(&err.to_string())
),
})?;
let model = tract_onnx::onnx()
.model_for_path(config.model_dir.join(MODEL_FILE))
.map_err(|err| SafetyNetError::ModelUnavailable {
reason: format!(
"failed to load kiji tract model: {}",
sanitize_error(&err.to_string())
),
})?;
let has_token_type_ids = model
.input_outlets()
.map_err(|err| SafetyNetError::ModelUnavailable {
reason: format!(
"failed to inspect kiji tract model: {}",
sanitize_error(&err.to_string())
),
})?
.iter()
.any(|outlet| model.node(outlet.node).name == "token_type_ids");
let model = model
.into_optimized()
.map_err(|err| SafetyNetError::ModelUnavailable {
reason: format!(
"failed to optimize kiji tract model: {}",
sanitize_error(&err.to_string())
),
})?;
let model = model
.into_runnable()
.map_err(|err| SafetyNetError::ModelUnavailable {
reason: format!(
"failed to prepare kiji tract model: {}",
sanitize_error(&err.to_string())
),
})?;
Ok(Self {
config,
tokenizer,
model: Mutex::new(model),
has_token_type_ids,
})
}
}
impl KijiDistilbertBackend for TractKijiBackend {
fn id(&self) -> &str {
"kiji-distilbert-tract"
}
fn version(&self) -> &str {
&self.config.version
}
fn decoding_params(&self) -> &[(&str, String)] {
&self.config.decoding_params
}
fn infer(&self, clean: &str) -> Result<Vec<RawSpan>, SafetyNetError> {
let actual = clean.len();
if actual > self.config.max_input_bytes {
return Err(SafetyNetError::InputTooLarge {
limit: self.config.max_input_bytes,
actual,
});
}
let encoded =
self.tokenizer
.encode(clean, true)
.map_err(|err| SafetyNetError::Runtime {
message: format!(
"kiji tokenizer failed: {}",
sanitize_error(&err.to_string())
),
})?;
let offsets = encoded.get_offsets();
let ids = encoded.get_ids();
let attention = encoded.get_attention_mask();
if ids.is_empty() {
return Ok(Vec::new());
}
let seq_len = ids.len();
let shape = [1usize, seq_len];
let input_ids: Vec<i64> = ids.iter().map(|&value| value as i64).collect();
let attn_mask: Vec<i64> = attention.iter().map(|&value| value as i64).collect();
let input_ids_tensor = Tensor::from_shape(&shape, &input_ids)
.map_err(|err| SafetyNetError::Runtime {
message: format!(
"kiji tract input_ids tensor failed: {}",
sanitize_error(&err.to_string())
),
})?
.into_tvalue();
let attn_tensor = Tensor::from_shape(&shape, &attn_mask)
.map_err(|err| SafetyNetError::Runtime {
message: format!(
"kiji tract attention_mask tensor failed: {}",
sanitize_error(&err.to_string())
),
})?
.into_tvalue();
let inputs = if self.has_token_type_ids {
let token_type = vec![0i64; seq_len];
let type_tensor = Tensor::from_shape(&shape, &token_type)
.map_err(|err| SafetyNetError::Runtime {
message: format!(
"kiji tract token_type_ids tensor failed: {}",
sanitize_error(&err.to_string())
),
})?
.into_tvalue();
tvec!(input_ids_tensor, attn_tensor, type_tensor)
} else {
tvec!(input_ids_tensor, attn_tensor)
};
let model = self.model.lock().map_err(|err| SafetyNetError::Runtime {
message: format!(
"kiji tract model lock poisoned: {}",
sanitize_error(&err.to_string())
),
})?;
let outputs = model.run(inputs).map_err(|err| SafetyNetError::Runtime {
message: format!(
"kiji tract inference failed: {}",
sanitize_error(&err.to_string())
),
})?;
let Some(output) = outputs.first() else {
return Ok(Vec::new());
};
let tensor: &Tensor = output;
let shape = tensor.shape();
if shape.len() != 3 || shape[0] != 1 || shape[1] != seq_len {
return Err(SafetyNetError::InvalidOutput {
message: "kiji tract returned invalid logits shape".to_string(),
});
}
let num_labels = shape[2];
let flat = tensor
.try_as_plain()
.and_then(|view| view.as_slice::<f32>())
.map_err(|err| SafetyNetError::Runtime {
message: format!(
"kiji tract output slice failed: {}",
sanitize_error(&err.to_string())
),
})?;
normalize_raw_spans(
decode_logits(clean, offsets, flat, seq_len, num_labels),
clean,
)
}
}
fn sanitize_error(message: &str) -> String {
message
.split_ascii_whitespace()
.map(sanitize_token)
.collect::<Vec<_>>()
.join(" ")
}
fn sanitize_token(token: &str) -> String {
if token.contains('@') {
return "<redacted>".to_string();
}
let digit_count = token.bytes().filter(u8::is_ascii_digit).count();
if digit_count >= 7 {
return "<redacted>".to_string();
}
token.to_string()
}