xberg 1.1.5

High-performance document intelligence library for Rust. Extract text, metadata, and structured data from PDFs, Office documents, images, and 107 formats and 371 programming languages via tree-sitter code intelligence with async/sync APIs.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72

//! Engine-neutral tensor currency for the inference seam.
//!
//! [`InferenceTensor`] is the type exchanged across the [`InferenceSession`] API,
//! independent of the backing engine (ONNX Runtime today; tract on no-ORT targets
//! in a later phase). Each variant wraps an owned dynamic-dimensional
//! [`ndarray::ArrayD`] — callers build inputs with [`ndarray`] and read outputs
//! back as views or slices. Engine-specific conversions (`ort::Value` ↔ tensor)
//! live next to their backend, keeping this module pure `ndarray`.
//!
//! [`InferenceSession`]: super::InferenceSession
//!

use ndarray::ArrayD;

/// A dense tensor passed across the inference seam.
///
/// The dtype set mirrors the ONNX tensor element types xberg's models actually
/// exchange. New dtypes are added here as models that need them are migrated onto
/// the seam.
#[derive(Debug, Clone, PartialEq)]
pub enum InferenceTensor {
    /// 32-bit float — the common image/logit currency.
    F32(ArrayD<f32>),
    /// 64-bit signed integer — token ids, shape inputs.
    I64(ArrayD<i64>),
    /// 32-bit signed integer — box counts and similar.
    I32(ArrayD<i32>),
    /// 8-bit unsigned integer — raw pixel inputs.
    U8(ArrayD<u8>),
    /// Boolean — attention/pad masks that ship as bool tensors.
    Bool(ArrayD<bool>),
}

impl InferenceTensor {
    /// Borrow the `f32` payload, or `None` when the tensor holds another dtype.
    ///
    /// `mod tensor` compiles unconditionally (this type is the engine-neutral seam's shared
    /// currency), but its real callers each sit behind their own, mutually distinct feature
    /// gates — `layout::models::table_classifier` under `layout-detection`/`pdf`,
    /// `doc_orientation::detector` under `auto_rotate`, and `inference::tract_backend` under its
    /// `tract`/`inference_ort` combination — so no single narrower `cfg` on this method covers
    /// exactly its live callers without also going dead in some other feature slice. Kept as an
    /// `allow` rather than a cfg for that reason; the `#[cfg(test)]` unit tests below exercise it
    /// unconditionally regardless of which (if any) of those features are enabled.
    #[allow(dead_code)]
    pub fn as_f32(&self) -> Option<&ArrayD<f32>> {
        match self {
            Self::F32(array) => Some(array),
            _ => None,
        }
    }
}

impl From<ArrayD<f32>> for InferenceTensor {
    fn from(array: ArrayD<f32>) -> Self {
        Self::F32(array)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn as_f32_returns_payload_for_f32_and_none_otherwise() {
        let f32_tensor: InferenceTensor = ndarray::arr1(&[1.0f32, 2.0]).into_dyn().into();
        assert_eq!(f32_tensor.as_f32().unwrap(), &ndarray::arr1(&[1.0f32, 2.0]).into_dyn());

        let i64_tensor = InferenceTensor::I64(ndarray::arr1(&[1i64, 2]).into_dyn());
        assert!(i64_tensor.as_f32().is_none());
    }
}