mirror of
https://github.com/pykeio/ort
synced 2026-04-25 16:34:55 +02:00
35 lines
1.1 KiB
Rust
35 lines
1.1 KiB
Rust
#![cfg(not(target_arch = "aarch64"))]
|
|
|
|
use std::path::Path;
|
|
|
|
use ndarray::{ArrayD, IxDyn};
|
|
use ort::{
|
|
inputs,
|
|
session::{Session, builder::GraphOptimizationLevel},
|
|
value::Tensor
|
|
};
|
|
|
|
#[test]
|
|
fn vectorizer() -> ort::Result<()> {
|
|
let mut session = Session::builder()?
|
|
.with_optimization_level(GraphOptimizationLevel::Level1)?
|
|
.with_intra_threads(1)?
|
|
.commit_from_file(Path::new(env!("CARGO_MANIFEST_DIR")).join("tests").join("data").join("vectorizer.onnx"))
|
|
.expect("Could not load model");
|
|
|
|
{
|
|
let metadata = session.metadata()?;
|
|
assert_eq!(metadata.producer()?, "skl2onnx");
|
|
assert_eq!(metadata.description()?, "test description");
|
|
assert_eq!(metadata.custom_keys()?, ["custom_key"]);
|
|
assert_eq!(metadata.custom("custom_key")?.as_deref(), Some("custom_value"));
|
|
}
|
|
|
|
let array = ndarray::CowArray::from(ndarray::Array::from_shape_vec((1,), vec!["document".to_owned()]).unwrap());
|
|
|
|
let outputs = session.run(inputs![Tensor::from_string_array(&array)?])?;
|
|
assert_eq!(outputs[0].try_extract_tensor::<f32>()?, ArrayD::from_shape_vec(IxDyn(&[1, 9]), vec![0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]).unwrap());
|
|
|
|
Ok(())
|
|
}
|