| 148 | |
| 149 | |
| 150 | class TFEncoderLayer(tf.keras.layers.Layer): |
| 151 | def __init__(self, d_model_size, num_heads, dff, rate=0.1, layer_norm_epsilon=1e-6, **kwargs): |
| 152 | super().__init__(**kwargs) |
| 153 | |
| 154 | self.multi_head_attention = TFMultiHeadAttention(d_model_size, num_heads, name="multi_head_attention") |
| 155 | self.ffn = point_wise_feed_forward_network(d_model_size, dff, name="ffn") |
| 156 | |
| 157 | self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=layer_norm_epsilon, name="layernorm1") |
| 158 | self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=layer_norm_epsilon, name="layernorm2") |
| 159 | |
| 160 | self.dropout1 = tf.keras.layers.Dropout(rate) |
| 161 | self.dropout2 = tf.keras.layers.Dropout(rate) |
| 162 | |
| 163 | def call(self, inputs, training=False): |
| 164 | x, mask, layer_past, attention_mask, head_mask, use_cache, output_attentions = inputs |
| 165 | normed = self.layernorm1(x) |
| 166 | attn_outputs = self.multi_head_attention( |
| 167 | [normed, normed, normed, mask, layer_past, attention_mask, head_mask, use_cache, output_attentions], |
| 168 | training=training, |
| 169 | ) |
| 170 | attn_output = attn_outputs[0] |
| 171 | attn_output = self.dropout1(attn_output, training=training) |
| 172 | out1 = x + attn_output |
| 173 | |
| 174 | out2 = self.layernorm2(out1) |
| 175 | ffn_output = self.ffn(out2) |
| 176 | ffn_output = self.dropout2(ffn_output, training=training) |
| 177 | out2 = out1 + ffn_output |
| 178 | |
| 179 | outputs = (out2,) + attn_outputs[1:] |
| 180 | return outputs |
| 181 | |
| 182 | |
| 183 | @keras_serializable |