/**
|
* @license
|
* Copyright 2023 Google LLC.
|
* 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.
|
* =============================================================================
|
*/
|
/// <amd-module name="@tensorflow/tfjs-layers/dist/layers/nlp/modeling/transformer_decoder" />
|
/**
|
* Transformer decoder block implementation based on TFJS `Layer`.
|
*/
|
import { Tensor, serialization } from '@tensorflow/tfjs-core';
|
import { Activation } from '../../../activations';
|
import { Layer, LayerArgs, SymbolicTensor } from '../../../engine/topology';
|
import { Initializer, InitializerIdentifier } from '../../../initializers';
|
import { ActivationIdentifier } from '../../../keras_format/activation_config';
|
import { Shape } from '../../../keras_format/common';
|
import { Dense, Dropout } from '../../core';
|
import { LayerNormalization } from '../../normalization';
|
import { CachedMultiHeadAttention } from './cached_multihead_attention';
|
export declare interface TransformerDecoderArgs extends LayerArgs {
|
/**
|
* Integer. The hidden size of feedforward network.
|
*/
|
intermediateDim: number;
|
/**
|
* Integer. The number of heads in MultiHeadAttention.
|
*/
|
numHeads: number;
|
/**
|
* The dropout value, shared by MultiHeadAttention and feedforward network.
|
* Defaults to `0.`.
|
*/
|
dropout?: number;
|
/**
|
* The activation function of feedforward network.
|
* Defaults to `"relu"`.
|
*/
|
activation?: Activation | ActivationIdentifier;
|
/**
|
* The eps value in layer normalization components.
|
* Defaults to `1e-5`.
|
*/
|
layerNormEpsilon?: number;
|
/**
|
* The kernel initializer for the dense and multiheaded attention layers.
|
* Defaults to `"glorotUniform"`.
|
*/
|
kernelInitializer?: Initializer | InitializerIdentifier;
|
/**
|
* The bias initializer for the dense and multiheaded attention layers.
|
* Defaults to `"zeros"`.
|
*/
|
biasInitializer?: Initializer | InitializerIdentifier;
|
/**
|
* If true, the inputs to the attention layer(s) and the intermediate dense
|
* layer are normalized (similar to GPT-2). If set to false, outputs of
|
* attention layer and intermediate dense layer are normalized
|
* (similar to BERT).
|
* Defaults to `false`.
|
*/
|
normalizeFirst?: boolean;
|
}
|
export declare interface TransformerDecoderOptions {
|
/**
|
* decoderSequence: The decode input sequence.
|
*/
|
/**
|
* The encoder input sequence. For decoder only models (like GPT2), this
|
* should be left `null`. Once the model is called without an encoderSequence,
|
* you cannot call it again with encoderSequence.
|
*/
|
encoderSequence?: Tensor | SymbolicTensor;
|
/**
|
* A boolean Tensor, the padding mask of decoder sequence, must be of shape
|
* `[batchSize, decoderSequenceLength]`.
|
*/
|
decoderPaddingMask?: Tensor | SymbolicTensor;
|
/**
|
* A boolean Tensor. Customized decoder sequence mask, must be of shape
|
* `[batchSize, decoderSequenceLength, decoderSequenceLength]`.
|
*/
|
decoderAttentionMask?: Tensor;
|
/**
|
* A boolean Tensor, the padding mask of encoder sequence, must be of shape
|
* `[batchSize, encoderSequenceLength]`.
|
*/
|
encoderPaddingMask?: Tensor;
|
/**
|
* A boolean Tensor. Customized encoder sequence mask, must be of shape
|
* `[batchSize, encoderSequenceLength, encoderSequenceLength]`.
|
*/
|
encoderAttentionMask?: Tensor;
|
/**
|
* A dense float Tensor. The cache of key/values pairs in the self-attention
|
* layer. Has shape `[batchSize, 2, maxSeqLen, numHeads, keyDims]`.
|
*/
|
selfAttentionCache?: Tensor;
|
/**
|
* Integer or Integer Tensor. The index at which to update the
|
* `selfAttentionCache`. Usually, this is the index of the current token
|
* being processed during decoding.
|
*/
|
selfAttentionCacheUpdateIndex?: number;
|
/**
|
* A dense float Tensor. The cache of key/value pairs in the cross-attention
|
* layer. Has shape `[batchSize, 2, S, numHeads, keyDims]`.
|
*/
|
crossAttentionCache?: Tensor;
|
/**
|
* Integer or Integer Tensor. The index at which to update the
|
* `crossAttentionCache`. Usually, this is either `0` (compute the entire
|
* `crossAttentionCache`), or `null` (reuse a previously computed
|
* `crossAttentionCache`).
|
*/
|
crossAttentionCacheUpdateIndex?: number;
|
/**
|
* If true, a causal mask (masking out future input) is applied on the decoder
|
* sequence.
|
* Defaults to `true`.
|
*/
|
useCausalMask?: boolean;
|
}
|
/**
|
* Transformer decoder.
|
*
|
* This class follows the architecture of the transformer decoder layer in the
|
* paper [Attention is All You Need](https://arxiv.org/abs/1706.03762). Users
|
* can instantiate multiple instances of this class to stack up a decoder.
|
*
|
* By default, this layer will apply a causal mask to the decoder attention
|
* layer. This layer will correctly compute an attention mask from an implicit
|
* padding mask (for example, by passing `maskZero=true` to a
|
* `tf.layers.embedding` layer). See the Masking and Padding
|
* [guide](https://keras.io/guides/understanding_masking_and_padding/)
|
* for more details.
|
*
|
* This layer can be called with either one or two inputs. The number of inputs
|
* must be consistent across all calls. The options are as follows:
|
* `layer.call(decoderSequence)`: no cross-attention will be built into the
|
* decoder block. This is useful when building a "decoder-only"
|
* transformer such as GPT-2.
|
* `layer.call(decoderSequence, {encoderSequence})`: cross-attention will be
|
* built into the decoder block. This is useful when building an
|
* "encoder-decoder" transformer, such as the original transformer
|
* model described in Attention is All You Need.
|
*
|
* Examples:
|
* ```js
|
* // Create a single transformer decoder layer.
|
* const decoder = new TransformerDecoder({intermediateDim: 64, numHeads: 8});
|
*
|
* // Create a simple model containing the decoder.
|
* const decoderInput = tf.input({shape: [10, 64]});
|
* const encoderInput = tf.input({shape: {[10, 64]});
|
* const output = decoder.call(decoderInput, {encoderInput});
|
* const model = tf.model({
|
* inputs: [decoderInput, encoderInput],
|
* outputs: output,
|
* );
|
*
|
* // Call decoder on the inputs.
|
* const decoderInputData = tf.randomUniform([2, 10, 64]);
|
* const encoderInputData = tf.randomUniform([2, 10, 64]);
|
* const decoderOutput = model.predict([decoderInputData, encoderInputData]);
|
* ```
|
*
|
* References:
|
* - [Vaswani et al., 2017](https://arxiv.org/abs/1706.03762)
|
*/
|
export declare class TransformerDecoder extends Layer {
|
/** @nocollapse */
|
static readonly className = "TransformerDecoder";
|
protected intermediateDim: number;
|
protected numHeads: number;
|
protected dropout: number;
|
protected activation: Activation;
|
protected layerNormEpsilon: number;
|
protected kernelInitializer: Initializer;
|
protected biasInitializer: Initializer;
|
protected normalizeFirst: boolean;
|
protected decoderSequenceShape: Shape;
|
protected encoderSequenceShape: Shape;
|
protected selfAttentionLayer: CachedMultiHeadAttention;
|
protected selfAttentionLayernorm: LayerNormalization;
|
protected selfAttentionDropout: Dropout;
|
protected selfCrossAttentionLayer: CachedMultiHeadAttention;
|
protected selfCrossAttentionLayernorm: LayerNormalization;
|
protected selfCrossAttentionDropout: Dropout;
|
protected feedforwardIntermediateDense: Dense;
|
protected feedforwardOutputDense: Dense;
|
protected feedforwardLayernorm: LayerNormalization;
|
protected feedforwardDropout: Dropout;
|
constructor(args: TransformerDecoderArgs);
|
/**
|
*
|
* @param inputShape decoderSequenceShape or
|
* [decoderSequenceShape, encoderSequenceShape]
|
*/
|
build(inputShape: Shape | [Shape, Shape]): void;
|
apply(decoderSequence: Tensor | SymbolicTensor, kwargs?: TransformerDecoderOptions): Tensor | SymbolicTensor;
|
call(decoderSequence: Tensor, kwargs: TransformerDecoderOptions): Tensor;
|
/**
|
* Forward pass of the TransformerDecoder.
|
*
|
* @returns One of three things, depending on call arguments:
|
* - `[outputs, null, null]`, if `selfAttentionCache` is `null`.
|
* - `[outputs, selfAttentionCache, null]`, if `selfAttentionCache` is
|
* set and the layer has no cross-attention.
|
* - `[outputs, selfAttentionCache, crossAttentionCache]`, if
|
* `selfAttentionCache` and `crossAttentionCache` are set and
|
* the layer has cross-attention.
|
*/
|
callAndReturnCaches(decoderSequence: Tensor, kwargs: TransformerDecoderOptions): [Tensor, Tensor, Tensor];
|
private computeSelfAttentionMask;
|
getConfig(): serialization.ConfigDict;
|
computeOutputShape(decoderSequenceShape: Shape): Shape;
|
}
|
