encoderfile 0.6.2

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

use encoderfile::transforms::{
    DEFAULT_LIBS, EmbeddingTransform, Postprocessor, SequenceClassificationTransform,
    TokenClassificationTransform,
};
use ndarray::{Array2, Array3, Axis};
use ort::tensor::ArrayExtensions;

#[test]
fn test_l2_normalization() {
    let engine = EmbeddingTransform::new(
        DEFAULT_LIBS.to_vec(),
        Some(include_str!("../../../transforms/embedding/l2_normalize_embeddings.lua").to_string()),
    )
    .expect("Failed to create engine");

    let test_arr = Array3::<f32>::from_elem((8, 16, 36), 1.0);

    let norms = test_arr
        .map_axis(Axis(2), |v| (v.mapv(|x| x.powi(2)).sum()).sqrt())
        .insert_axis(Axis(2)); // shape: [batch_size, n_tokens, 1]

    // Avoid divide-by-zero by adding small epsilon
    let eps = 1e-12;
    let gold = &test_arr / &(&norms + eps);

    let test = engine.postprocess(test_arr).expect("Didn't read");

    assert_eq!(gold, test)
}

#[test]
fn test_softmax_sequence_cls() {
    let engine = SequenceClassificationTransform::new(
        DEFAULT_LIBS.to_vec(),
        Some(
            include_str!("../../../transforms/sequence_classification/softmax_logits.lua")
                .to_string(),
        ),
    )
    .expect("Failed to create engine");

    // run on array of shape [batch_size, n_labels]
    let test_arr = Array2::<f32>::from_elem((8, 16), 1.0);

    let softmax_gold = test_arr.softmax(Axis(1));

    let softmax_test = engine
        .postprocess(test_arr)
        .expect("Failed to compute softmax");

    assert_eq!(softmax_gold, softmax_test);
}

#[test]
fn test_softmax_token_cls() {
    let engine = TokenClassificationTransform::new(
        DEFAULT_LIBS.to_vec(),
        Some(
            include_str!("../../../transforms/token_classification/softmax_logits.lua").to_string(),
        ),
    )
    .expect("Failed to create engine");

    // run on array of shape [batch_size, n_tokens, n_labels]
    let test_arr = Array3::<f32>::from_elem((8, 16, 2), 1.0);

    let softmax_gold = test_arr.softmax(Axis(2));

    let softmax_test = engine
        .postprocess(test_arr)
        .expect("Failed to compute softmax");

    assert_eq!(softmax_gold, softmax_test);
}