/
githubmirror
/
models
Обзор
Документация
Войти
/
githubmirror
/
models
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
master
official/projects/const_cl/modeling/heads/transformer_decoder.py
343 строки
11 KB
A. Unique TensorFlower
No public description
09 фев 2026, 19:00
09 фев 2026, 19:00
799b0af
Код
Авторство
О чём код?
# Copyright 2026 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Definition for Transformer heads.""" from typing import Any, Mapping, Optional, Union, List, Sequence from absl import logging import tensorflow as tf, tf_keras def _get_shape(x: tf.Tensor): """Helper function to return shape of a given tensor.""" static = x.shape.as_list() dynamic = tf.shape(x) return [dynamic[i] if s is None else s for i, s in enumerate(static)] class DecoderUnit(tf_keras.layers.Layer): """Constructs the decoder MHA module used in Transformer layers.""" def __init__(self, num_channels: int, use_bias: bool, dropout_rate: float, activation: str, layer_norm_epsilon: float, **kwargs): super().__init__(**kwargs) self._num_channels = num_channels self._use_bias = use_bias self._dropout_rate = dropout_rate self._activation = activation self._layer_norm_epsilon = layer_norm_epsilon def build(self, input_shape: Union[tf.TensorShape, List[tf.TensorShape]]): """Builds the layer. Args: input_shape: the input shape for the keras tensor. """ # Query, key, and value mapping. self.layer_q = tf_keras.layers.Dense( self._num_channels, use_bias=self._use_bias, activation=None, name='query') self.layer_k = tf_keras.layers.Dense( self._num_channels, use_bias=self._use_bias, activation=None, name='key') self.layer_v = tf_keras.layers.Dense( self._num_channels, use_bias=self._use_bias, activation=None, name='value') self.dropout = tf_keras.layers.Dropout(self._dropout_rate) # Note here is a different behavior for contrib_layers.layer_norm and # tf_keras.layers.LayerNormalization, where by default, the former # calculates mean/variance across all axes except the first one # (batch axis), while the latter one computes statistics only on the last # axis. self.layer_norm = tf_keras.layers.LayerNormalization( epsilon=self._layer_norm_epsilon, name='layer_norm') self.ffn1 = tf_keras.layers.Dense( self._num_channels, use_bias=self._use_bias, activation=self._activation, name='ffn1') self.ffn2 = tf_keras.layers.Dense( self._num_channels, use_bias=self._use_bias, activation=None, name='ffn2') super().build(input_shape) def call(self, query: tf.Tensor, memory: Optional[tf.Tensor], training: bool = False) -> Mapping[str, tf.Tensor]: """Forward pass of the Transformer decoder unit. Args: query: the input query tensor. memory: the input memory tensor for key/value pairs. If None, self-attention will be performed. training: whether in training mode. Returns: outputs: the output dictionary contains 'hidden_states' and 'attention weights' matrix. """ if memory is None: memory = query tensor_q = self.layer_q(query) # (bs, qlen, inner_dim) tensor_k = self.layer_k(memory) # (bs, klen, inner_dim) tensor_v = self.layer_v(memory) # (bs, klen, inner_dim) scores = tf.matmul(tensor_q, tensor_k, transpose_b=True) # Scales attention_scores. dk = tf.cast(_get_shape(tensor_k)[-1], dtype=scores.dtype) scores = scores / tf.math.sqrt(dk) # Shape: (bs, seq_len, seq_len) attention_weights = tf.nn.softmax(scores, axis=-1) # Shape: (bs, seq_len, dim_per_head) attention_features = tf.matmul(attention_weights, tensor_v) # Shape: (bs, seq_len, seq_len) attention_features = self.dropout(attention_features, training=training) hidden_states = attention_features + tensor_q hidden_states = self.layer_norm(hidden_states) # Shape: (bs, seq_len, out_dim) hidden_states = self.ffn1(hidden_states) hidden_states = self.ffn2(hidden_states) outputs = { 'hidden_states': hidden_states, 'attention_weights': attention_weights, } return outputs def get_config(self) -> Mapping[str, Any]: """Gets class config parameters.""" config_dict = { 'num_channels': self._num_channels, 'use_bias': self._use_bias, 'dropout_rate': self._dropout_rate, 'activation': self._activation, 'layer_norm_epsilon': self._layer_norm_epsilon, } return config_dict @classmethod def from_config(cls, config: Mapping[str, Any]): """Factory constructor from config.""" return cls(**config) class TransformerDecoderLayer(tf_keras.layers.Layer): """Constructs the main Transformer decoder module which includes MHA + FFN.""" def __init__(self, num_channels: int, num_heads: int, use_bias: bool, activation: str, dropout_rate: float, layer_norm_epsilon: float, name: str = 'decoder_layer', **kwargs): super().__init__(name=name) self._num_channels = num_channels self._num_heads = num_heads self._use_bias = use_bias self._activation = activation self._dropout_rate = dropout_rate self._layer_norm_epsilon = layer_norm_epsilon self._name = name self._mha_units = [] for i in range(num_heads): self._mha_units.append( DecoderUnit( num_channels=num_channels, use_bias=use_bias, dropout_rate=dropout_rate, activation=activation, layer_norm_epsilon=layer_norm_epsilon, name='mha_{}'.format(i))) def call( self, inputs: tf.Tensor, memory: Optional[tf.Tensor] = None, training: bool = False ) -> Mapping[str, Union[tf.Tensor, Sequence[tf.Tensor]]]: """Forward pass of the Transformer decoder layer. Args: inputs: the input query tensor. memory: the input memory tensor for key/value pairs. If None, self-attention will be performed. training: whether in training mode. Returns: outputs: the output dictionary contains 'hidden_states' and 'attention weights' matrix. """ if memory is None: logging.info('No memory tokens are provided. Performing self-attention ' 'on input tokens in TransfomerDecoder.') all_head_feats = [] all_head_attentions = [] for i in range(self._num_heads): outputs = self._mha_units[i]( query=inputs, memory=memory, training=training) all_head_feats.append(outputs['hidden_states']) all_head_attentions.append(outputs['attention_weights']) outputs = { 'hidden_states': tf.concat(all_head_feats, axis=-1), 'attention_weights': all_head_attentions, } return outputs def get_config(self) -> Mapping[str, Any]: """Gets class config parameters.""" config_dict = { 'num_channels': self._num_channels, 'num_heads': self._num_heads, 'use_bias': self._use_bias, 'activation': self._activation, 'dropout_rate': self._dropout_rate, 'layer_norm_epsilon': self._layer_norm_epsilon, 'name': self._name, } return config_dict @classmethod def from_config(cls, config: Mapping[str, Any]): """Factory constructor from config.""" return cls(**config) class TransformerDecoder(tf_keras.layers.Layer): """Constructs the final Transformer decoder stack.""" def __init__(self, num_channels: int, num_layers: int, num_heads: int, use_bias: bool, activation: str, dropout_rate: float, layer_norm_epsilon: float, name: str = 'transformer_decoder', **kwargs): super().__init__(name=name) self._num_channels = num_channels self._num_layers = num_layers self._num_heads = num_heads self._use_bias = use_bias self._activation = activation self._dropout_rate = dropout_rate self._layer_norm_epsilon = layer_norm_epsilon self._layers = [] for n in range(self._num_layers): self._layers.append( TransformerDecoderLayer( num_channels=num_channels, num_heads=num_heads, use_bias=use_bias, activation=activation, dropout_rate=dropout_rate, layer_norm_epsilon=layer_norm_epsilon, name='layer_{}'.format(n))) def call(self, inputs: tf.Tensor, memory: Optional[tf.Tensor] = None, training: bool = False) -> Mapping[str, Sequence[tf.Tensor]]: """Forward pass of the Transformer decoder. Args: inputs: the input query tensor. memory: the input memory tensor for key/value pairs. If None, self-attention will be performed. training: whether in training mode. Returns: outputs: the output dictionary contains 'hidden_states' and 'attention weights' matrix. """ all_hidden_states = () all_attentions = () memory_shape = _get_shape(memory) memory = tf.reshape(memory, [memory_shape[0], -1, memory_shape[-1]]) hidden_states = inputs for layer in self._layers: layer_outputs = layer(inputs=hidden_states, memory=memory, training=training) # layer_outputs is a dictionary with the following keys: # hidden_states, self_attention_weights hidden_states = layer_outputs['hidden_states'] all_attentions += (layer_outputs['attention_weights'],) # Add last layer all_hidden_states += (hidden_states,) outputs = { 'hidden_states': all_hidden_states, 'attention_weights': all_attentions, } return outputs def get_config(self) -> Mapping[str, Any]: """Gets class config parameters.""" config_dict = { 'num_channels': self._num_channels, 'num_layers': self._num_layers, 'num_heads': self._num_heads, 'use_bias': self._use_bias, 'activation': self._activation, 'dropout_rate': self._dropout_rate, 'layer_norm_epsilon': self._layer_norm_epsilon, } return config_dict @classmethod def from_config(cls, config: Mapping[str, Any]): """Factory constructor from config.""" return cls(**config)