1*a5412caaSRahul use anyhow::{Context, Result};
2*a5412caaSRahul use std::fs;
3*a5412caaSRahul use test_programs::nn::{sort_results, witx};
4*a5412caaSRahul 
main() -> Result<()>5*a5412caaSRahul pub fn main() -> Result<()> {
6*a5412caaSRahul     let model = fs::read("fixture/model.pt")
7*a5412caaSRahul         .context("the model file to be mapped to the fixture directory")?;
8*a5412caaSRahul     let graph = witx::load(
9*a5412caaSRahul         &[&model],
10*a5412caaSRahul         witx::GraphEncoding::Pytorch,
11*a5412caaSRahul         witx::ExecutionTarget::CPU,
12*a5412caaSRahul     )?;
13*a5412caaSRahul     let tensor = fs::read("fixture/kitten.tensor")
14*a5412caaSRahul         .context("the tensor file to be mapped to the fixture directory")?;
15*a5412caaSRahul     let output_buffer = witx::classify(graph, tensor)?;
16*a5412caaSRahul     let result = softmax(output_buffer);
17*a5412caaSRahul     let top_five = &sort_results(&result)[..5];
18*a5412caaSRahul     assert_eq!(top_five[0].class_id(), 281);
19*a5412caaSRahul     println!("found results, sorted top 5: {top_five:?}");
20*a5412caaSRahul     Ok(())
21*a5412caaSRahul }
22*a5412caaSRahul 
softmax(output_tensor: Vec<f32>) -> Vec<f32>23*a5412caaSRahul fn softmax(output_tensor: Vec<f32>) -> Vec<f32> {
24*a5412caaSRahul     let max_val = output_tensor
25*a5412caaSRahul         .iter()
26*a5412caaSRahul         .cloned()
27*a5412caaSRahul         .fold(f32::NEG_INFINITY, f32::max);
28*a5412caaSRahul 
29*a5412caaSRahul     // Compute the exponential of each element subtracted by max_val for numerical stability.
30*a5412caaSRahul     let exps: Vec<f32> = output_tensor.iter().map(|&x| (x - max_val).exp()).collect();
31*a5412caaSRahul 
32*a5412caaSRahul     // Compute the sum of the exponentials.
33*a5412caaSRahul     let sum_exps: f32 = exps.iter().sum();
34*a5412caaSRahul 
35*a5412caaSRahul     // Normalize each element to get the probabilities.
36*a5412caaSRahul     exps.iter().map(|&exp| exp / sum_exps).collect()
37*a5412caaSRahul }
38