mirror of
https://github.com/pykeio/ort
synced 2026-04-25 16:34:55 +02:00
d269461e21
The only error case occurs when the value is not present, so `Option<T>`s make more sense here.
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().as_deref(), Some("skl2onnx"));
|
|
assert_eq!(metadata.description().as_deref(), Some("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_array::<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(())
|
|
}
|