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
88
89
90
91

use crate::TransformerError;
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,
        }
    }

    /// Map the output representation of this layer.
    pub fn map_output<F>(&self, f: F) -> Result<Self, TransformerError>
    where
        F: Fn(&Tensor) -> Result<Tensor, TransformerError>,
    {
        let layer = match self {
            LayerOutput::Embedding(embedding) => LayerOutput::Embedding(f(embedding)?),
            LayerOutput::EncoderWithAttention(HiddenLayer { output, attention }) => {
                LayerOutput::EncoderWithAttention(HiddenLayer {
                    output: f(output)?,
                    attention: attention.shallow_clone(),
                })
            }
        };

        Ok(layer)
    }

    /// 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),
        }
    }
}