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
use tch::Tensor;
/// Hidden layer output and attention.
#[derive(Debug)]
pub struct HiddenLayer {
/// The output of the layer.
pub output: Tensor,
/// The layer attention scores (unnormalized).
pub attention: Tensor,
}
/// Output of a BERT layer.
#[derive(Debug)]
pub enum LayerOutput {
/// Embedding layer output.
Embedding(Tensor),
/// Encoder layer output.
EncoderWithAttention(HiddenLayer),
}
impl LayerOutput {
/// Get the layer attention.
///
/// Return a `Some` value if the layer output is from an encoder layer,
/// or `None` otherwise.
pub fn attention(&self) -> Option<&Tensor> {
match self {
LayerOutput::Embedding(_) => None,
LayerOutput::EncoderWithAttention(hidden) => Some(&hidden.attention),
}
}
/// Get the embedding.
///
/// Returns `Some` if the layer output is an embedding or `None`
/// otherwise.
pub fn embedding(&self) -> Option<&Tensor> {
match self {
LayerOutput::Embedding(embedding) => Some(embedding),
LayerOutput::EncoderWithAttention(_) => None,
}
}
/// Get the layer output.
pub fn output(&self) -> &Tensor {
match self {
LayerOutput::Embedding(embedding) => embedding,
LayerOutput::EncoderWithAttention(hidden) => &hidden.output,
}
}
/// Get the layer output mutably.
pub fn output_mut(&mut self) -> &mut Tensor {
match self {
LayerOutput::Embedding(embedding) => embedding,
LayerOutput::EncoderWithAttention(hidden) => &mut hidden.output,
}
}
/// Get the output of an encoder layer.
///
/// Return a `Some` value if the layer output is from an encoder layer,
/// or `None` otherwise.
pub fn encoder_with_attention(&self) -> Option<&HiddenLayer> {
match self {
LayerOutput::Embedding(_) => None,
LayerOutput::EncoderWithAttention(hidden) => Some(hidden),
}
}
}