Struct InferenceEngine 

pub struct InferenceEngine { /* private fields */ }

Inference engine for running ONNX model inference

Features

• Synchronous and asynchronous inference
• Binary and multi-label classification
• Automatic post-processing (softmax/sigmoid)
• Batch inference support (optional)

Example

use llm_shield_models::InferenceEngine;
use std::sync::Arc;

let engine = InferenceEngine::new(session);

// Run inference
let result = engine.infer(
    &input_ids,
    &attention_mask,
    &labels,
    PostProcessing::Softmax,
).await?;

println!("Predicted: {}", result.predicted_label().unwrap());
println!("Confidence: {:.2}", result.max_score);

Implementations

impl InferenceEngine

pub fn new(session: Arc<Mutex<Session>>) -> Self

Create a new inference engine

Arguments

• session - ONNX Runtime session wrapped in Arc<Mutex<>> for thread-safe mutable access

pub async fn infer_async( &self, input_ids: &[u32], attention_mask: &[u32], labels: &[String], post_processing: PostProcessing, ) -> Result<InferenceResult>

Run inference on input IDs (async)

Arguments

• input_ids - Tokenized input IDs
• attention_mask - Attention mask (1 for real tokens, 0 for padding)
• labels - Class labels
• post_processing - Post-processing method (Softmax or Sigmoid)

Returns

InferenceResult with predictions and confidence scores

pub fn infer( &self, input_ids: &[u32], attention_mask: &[u32], labels: &[String], post_processing: PostProcessing, ) -> Result<InferenceResult>

Run inference on input IDs (synchronous)

Arguments

• input_ids - Tokenized input IDs
• attention_mask - Attention mask (1 for real tokens, 0 for padding)
• labels - Class labels
• post_processing - Post-processing method (Softmax or Sigmoid)

Returns

InferenceResult with predictions and confidence scores

pub fn softmax_static(logits: &[f32]) -> Vec<f32>

Apply softmax to logits (static method)

Softmax converts logits to probabilities that sum to 1.0. Used for single-label classification (mutually exclusive classes).

Arguments

• logits - Raw model output logits

Returns

Probability distribution (sums to 1.0)

Example

use llm_shield_models::InferenceEngine;

let logits = vec![1.0, 2.0, 0.5];
let probs = InferenceEngine::softmax_static(&logits);

// Probabilities sum to 1.0
let sum: f32 = probs.iter().sum();
assert!((sum - 1.0).abs() < 0.001);

pub fn sigmoid_static(logits: &[f32]) -> Vec<f32>

Apply sigmoid to logits (static method)

Sigmoid converts each logit independently to [0, 1]. Used for multi-label classification (non-exclusive classes).

Arguments

• logits - Raw model output logits

Returns

Independent probabilities (do NOT sum to 1.0)

Example

use llm_shield_models::InferenceEngine;

let logits = vec![0.0, 2.0, -2.0];
let probs = InferenceEngine::sigmoid_static(&logits);

// sigmoid(0) ≈ 0.5
assert!((probs[0] - 0.5).abs() < 0.01);

// All probabilities in [0, 1]
for p in probs {
    assert!(p >= 0.0 && p <= 1.0);
}

pub async fn infer_token_classification( &self, input_ids: &[u32], attention_mask: &[u32], labels: &[String], ) -> Result<Vec<TokenPrediction>>

Run token-level classification inference (for NER/token classification)

Arguments

• input_ids - Token IDs from tokenizer
• attention_mask - Attention mask (1=real token, 0=padding)
• labels - BIO tag labels (e.g., [“O”, “B-PERSON”, “I-PERSON”, …])

Returns

Vector of predictions, one per input token

Errors

• Model inference failure
• Invalid tensor shapes
• Label count mismatch

Example

use llm_shield_models::InferenceEngine;

let engine = InferenceEngine::new(session);
let labels = vec!["O", "B-PERSON", "I-PERSON"];

let predictions = engine.infer_token_classification(
    &input_ids,
    &attention_mask,
    &labels
).await?;

for pred in predictions {
    println!("{}: {:.2}", pred.predicted_label, pred.confidence);
}