rust-bert 0.23.0

Ready-to-use NLP pipelines and language models
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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87

use crate::RustBertError;
use ndarray::{ArrayBase, ArrayD, CowArray, CowRepr, IxDyn};

use ort::{Session, Value};
use std::convert::{TryFrom, TryInto};
use tch::{Kind, Tensor};

pub(crate) fn ort_tensor_to_tch(ort_tensor: &Value) -> Result<Tensor, RustBertError> {
    let ort_tensor = ort_tensor.try_extract::<f32>()?.view().to_owned();
    Ok(Tensor::try_from(ort_tensor)?)
}

pub(crate) fn array_to_ort<'a>(
    session: &Session,
    array: &'a TypedArray<'a>,
) -> Result<Value<'a>, RustBertError> {
    match &array {
        TypedArray::I64(array) => Ok(Value::from_array(session.allocator(), array)?),
        TypedArray::F32(array) => Ok(Value::from_array(session.allocator(), array)?),
        TypedArray::I32(array) => Ok(Value::from_array(session.allocator(), array)?),
        TypedArray::F64(array) => Ok(Value::from_array(session.allocator(), array)?),
        TypedArray::F16(array) => Ok(Value::from_array(session.allocator(), array)?),
        TypedArray::I16(array) => Ok(Value::from_array(session.allocator(), array)?),
        TypedArray::I8(array) => Ok(Value::from_array(session.allocator(), array)?),
        TypedArray::UI8(array) => Ok(Value::from_array(session.allocator(), array)?),
        TypedArray::BF16(array) => Ok(Value::from_array(session.allocator(), array)?),
    }
}

pub(crate) enum TypedArray<'a> {
    I64(ArrayBase<CowRepr<'a, i64>, IxDyn>),
    F32(ArrayBase<CowRepr<'a, f32>, IxDyn>),
    I32(ArrayBase<CowRepr<'a, i32>, IxDyn>),
    F64(ArrayBase<CowRepr<'a, f64>, IxDyn>),
    F16(ArrayBase<CowRepr<'a, half::f16>, IxDyn>),
    I16(ArrayBase<CowRepr<'a, i16>, IxDyn>),
    I8(ArrayBase<CowRepr<'a, i8>, IxDyn>),
    UI8(ArrayBase<CowRepr<'a, u8>, IxDyn>),
    BF16(ArrayBase<CowRepr<'a, half::bf16>, IxDyn>),
}

pub(crate) fn tch_tensor_to_ndarray(tch_tensor: &Tensor) -> Result<TypedArray, RustBertError> {
    let kind = tch_tensor.kind();
    Ok(match kind {
        Kind::Int64 => {
            let array: ArrayD<i64> = tch_tensor.try_into()?;
            TypedArray::I64(CowArray::from(array))
        }
        Kind::Float => {
            let array: ArrayD<f32> = tch_tensor.try_into()?;
            TypedArray::F32(CowArray::from(array))
        }
        Kind::Int => {
            let array: ArrayD<i32> = tch_tensor.try_into()?;
            TypedArray::I32(CowArray::from(array))
        }
        Kind::Double => {
            let array: ArrayD<f64> = tch_tensor.try_into()?;
            TypedArray::F64(CowArray::from(array))
        }
        Kind::Half => {
            let array: ArrayD<half::f16> = tch_tensor.try_into()?;
            TypedArray::F16(CowArray::from(array))
        }
        Kind::Int16 => {
            let array: ArrayD<i16> = tch_tensor.try_into()?;
            TypedArray::I16(CowArray::from(array))
        }
        Kind::Int8 => {
            let array: ArrayD<i8> = tch_tensor.try_into()?;
            TypedArray::I8(CowArray::from(array))
        }
        Kind::Uint8 => {
            let array: ArrayD<u8> = tch_tensor.try_into()?;
            TypedArray::UI8(CowArray::from(array))
        }
        Kind::BFloat16 => {
            let array: ArrayD<half::bf16> = tch_tensor.try_into()?;
            TypedArray::BF16(CowArray::from(array))
        }
        _ => {
            return Err(RustBertError::ValueError(format!(
                "Type not supported: attempted to get convert torch tensor to ndarray for {kind:?}",
            )))
        }
    })
}