encoderfile 0.4.0-rc.1

Distribute and run transformer encoders with a single file.
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

use ndarray::{Array2, Axis};
use ort::session::Session;
use parking_lot::MutexGuard;

#[macro_export]
macro_rules! prepare_inputs {
    ($encodings:ident) => {{
        let padded_token_length = $encodings[0].len();

        // Get token IDs & mask as a flattened array.
        let ids: Vec<i64> = $encodings
            .iter()
            .flat_map(|e| e.get_ids().iter().map(|i| *i as i64))
            .collect();
        let mask: Vec<i64> = $encodings
            .iter()
            .flat_map(|e| e.get_attention_mask().iter().map(|i| *i as i64))
            .collect();
        let type_ids: Vec<i64> = $encodings
            .iter()
            .flat_map(|e| e.get_type_ids().iter().map(|i| *i as i64))
            .collect();

        // Convert flattened arrays into 2-dimensional tensors of shape [N, L].
        let a_ids = ort::value::TensorRef::from_array_view((
            [$encodings.len(), padded_token_length],
            &*ids,
        ))
        .unwrap()
        .to_owned();
        let a_mask = ort::value::TensorRef::from_array_view((
            [$encodings.len(), padded_token_length],
            &*mask,
        ))
        .unwrap()
        .to_owned();
        let a_type_ids = ort::value::TensorRef::from_array_view((
            [$encodings.len(), padded_token_length],
            &*type_ids,
        ))
        .unwrap()
        .to_owned();

        (a_ids, a_mask, a_type_ids)
    }};
}

pub fn l2_normalize(mut x: Array2<f32>, axis: Axis) -> Array2<f32> {
    for mut row in x.axis_iter_mut(axis) {
        let norm = row.mapv(|v| v * v).sum().sqrt();
        if norm > 0.0 {
            row.mapv_inplace(|v| v / norm);
        }
    }
    x
}

pub fn requires_token_type_ids<'a>(session: &MutexGuard<'a, Session>) -> bool {
    session
        .inputs
        .iter()
        .any(|inp| inp.name == "token_type_ids")
}

#[macro_export]
macro_rules! run_model {
    ($session:expr, $a_ids:expr, $a_mask:expr, $a_type_ids:expr) => {{
        match $crate::inference::utils::requires_token_type_ids(&$session) {
            true => $session.run(ort::inputs!($a_ids, $a_mask, $a_type_ids)),
            false => $session.run(ort::inputs!($a_ids, $a_mask)),
        }
        .map_err(|e| {
            tracing::error!("Error running model: {:?}", e);
            $crate::error::ApiError::InternalError("Error running model")
        })
    }};
}