uni_xervo/traits/
raw_tensor_model.rs

1use crate::error::Result;
2use crate::traits::ModelInfo;
3use async_trait::async_trait;
4use indexmap::IndexMap;
5use ndarray::ArrayD;
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
8pub enum TensorDtype {
9    F32,
10    F64,
11    I32,
12    I64,
13    Bool,
14    String,
15}
16
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub enum DimSize {
19    Fixed(usize),
20    Dynamic,
21    Named(String),
22}
23
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct TensorSpec {
26    pub name: String,
27    pub dtype: TensorDtype,
28    pub shape: Vec<DimSize>,
29}
30
31#[derive(Debug, Clone, PartialEq)]
32pub enum TensorValue {
33    F32(ArrayD<f32>),
34    F64(ArrayD<f64>),
35    I32(ArrayD<i32>),
36    I64(ArrayD<i64>),
37    Bool(ArrayD<bool>),
38    String(ArrayD<String>),
39}
40
41impl TensorValue {
42    pub fn dtype(&self) -> TensorDtype {
43        match self {
44            Self::F32(_) => TensorDtype::F32,
45            Self::F64(_) => TensorDtype::F64,
46            Self::I32(_) => TensorDtype::I32,
47            Self::I64(_) => TensorDtype::I64,
48            Self::Bool(_) => TensorDtype::Bool,
49            Self::String(_) => TensorDtype::String,
50        }
51    }
52
53    pub fn shape(&self) -> &[usize] {
54        match self {
55            Self::F32(v) => v.shape(),
56            Self::F64(v) => v.shape(),
57            Self::I32(v) => v.shape(),
58            Self::I64(v) => v.shape(),
59            Self::Bool(v) => v.shape(),
60            Self::String(v) => v.shape(),
61        }
62    }
63
64    pub fn ndim(&self) -> usize {
65        self.shape().len()
66    }
67}
68
69#[derive(Debug, Clone, Default, PartialEq)]
70pub struct TensorBatch {
71    pub tensors: IndexMap<String, TensorValue>,
72}
73
74impl TensorBatch {
75    pub fn new() -> Self {
76        Self::default()
77    }
78
79    pub fn insert(&mut self, name: impl Into<String>, value: TensorValue) -> Option<TensorValue> {
80        self.tensors.insert(name.into(), value)
81    }
82
83    pub fn get(&self, name: &str) -> Option<&TensorValue> {
84        self.tensors.get(name)
85    }
86
87    pub fn into_tensors(self) -> IndexMap<String, TensorValue> {
88        self.tensors
89    }
90
91    pub fn len(&self) -> usize {
92        self.tensors.len()
93    }
94
95    pub fn is_empty(&self) -> bool {
96        self.tensors.is_empty()
97    }
98}
99
100#[async_trait]
101pub trait RawTensorModel: ModelInfo {
102    async fn run(&self, inputs: &TensorBatch) -> Result<TensorBatch>;
103
104    async fn run_batch(&self, inputs: &[TensorBatch]) -> Result<Vec<TensorBatch>> {
105        let mut out = Vec::with_capacity(inputs.len());
106        for input in inputs {
107            out.push(self.run(input).await?);
108        }
109        Ok(out)
110    }
111
112    fn max_batch_size(&self) -> usize {
113        1
114    }
115
116    fn input_signature(&self) -> &[TensorSpec];
117    fn output_signature(&self) -> &[TensorSpec];
118
119    async fn warmup(&self) -> Result<()> {
120        Ok(())
121    }
122}
uni_xervo/traits/raw_tensor_model.rs

uni_xervo/traits/
raw_tensor_model.rs
