ai.djl.nn.transformer

Class BertBlock

java.lang.Object

ai.djl.nn.AbstractBlock

ai.djl.nn.transformer.BertBlock

All Implemented Interfaces:

Block

public final class BertBlock
extends AbstractBlock

Implements the core bert model (without next sentence and masked language task) of bert.

This closely follows the original Devlin et. al. paper and its reference implementation.

Nested Class Summary

Nested Classes

Modifier and Type	Class and Description
static class	BertBlock.Builder The Builder to construct a BertBlock type of Block.

Field Summary

Fields inherited from class ai.djl.nn.AbstractBlock

children, inputNames, inputShapes, parameters, parameterShapeCallbacks, version

Method Summary

Modifier and Type	Method and Description
static BertBlock.Builder	builder() Returns a new BertBlock builder.
static NDArray	createAttentionMaskFromInputMask(NDArray ids, NDArray mask) Creates a 3D attention mask from a 2D tensor mask.
NDList	forward(ParameterStore ps, NDArray tokenIds, NDArray typeIds, NDArray masks, boolean training) Embeds the input, runs it through a transformer encoder stack and outputs the resulting embeddings and the pooled classifier.
NDList	forward(ParameterStore ps, NDList inputs, boolean training) Applies the operating function of the block once.
protected NDList	forwardInternal(ParameterStore ps, NDList inputs, boolean training, ai.djl.util.PairList<java.lang.String,java.lang.Object> params)
int	getEmbeddingSize() Returns the embedding size used for tokens.
Shape[]	getOutputShapes(NDManager manager, Shape[] inputShapes) Returns the expected output shapes of the block for the specified input shapes.
int	getTokenDictionarySize() Returns the size of the token dictionary.
IdEmbedding	getTokenEmbedding() Returns the token embedding used by this Bert model.
int	getTypeDictionarySize() Returns the size of the type dictionary.
void	initializeChildBlocks(NDManager manager, DataType dataType, Shape... inputShapes) Initializes the Child blocks of this block.

Methods inherited from class ai.djl.nn.AbstractBlock

addChildBlock, addParameter, addParameter, addParameter, beforeInitialize, cast, clear, describeInput, forward, getChildren, getDirectParameters, getParameters, getParameterShape, initialize, isInitialized, loadMetadata, loadParameters, readInputShapes, saveInputShapes, saveMetadata, saveParameters, setInitializer, setInitializer, toString

Methods inherited from class java.lang.Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, wait, wait, wait

Methods inherited from interface ai.djl.nn.Block

forward, validateLayout

Method Detail

getTokenEmbedding

public IdEmbedding getTokenEmbedding()

Returns the token embedding used by this Bert model.

Returns:

the token embedding used by this Bert model

getEmbeddingSize

public int getEmbeddingSize()

Returns the embedding size used for tokens.

Returns:

the embedding size used for tokens

getTokenDictionarySize

public int getTokenDictionarySize()

Returns the size of the token dictionary.

Returns:

the size of the token dictionary

getTypeDictionarySize

public int getTypeDictionarySize()

Returns the size of the type dictionary.

Returns:

the size of the type dictionary

getOutputShapes

public Shape[] getOutputShapes(NDManager manager,
                               Shape[] inputShapes)

Returns the expected output shapes of the block for the specified input shapes.

Parameters:

manager - an NDManager

inputShapes - the shapes of the inputs

Returns:

the expected output shapes of the block

initializeChildBlocks

public void initializeChildBlocks(NDManager manager,
                                  DataType dataType,
                                  Shape... inputShapes)

Initializes the Child blocks of this block. You need to override this method if your subclass has child blocks. Used to determine the correct input shapes for child blocks based on the requested input shape for this block.

Overrides:

initializeChildBlocks in class AbstractBlock

Parameters:

manager - the manager to use for initialization

dataType - the requested data type

inputShapes - the expected input shapes for this block

createAttentionMaskFromInputMask

public static NDArray createAttentionMaskFromInputMask(NDArray ids,
                                                       NDArray mask)

Creates a 3D attention mask from a 2D tensor mask.

Parameters:

ids - 2D Tensor of shape (B, F)

mask - 2D Tensor of shape (B, T)

Returns:

float tensor of shape (B, F, T)

forwardInternal

protected NDList forwardInternal(ParameterStore ps,
                                 NDList inputs,
                                 boolean training,
                                 ai.djl.util.PairList<java.lang.String,java.lang.Object> params)

Specified by:

forwardInternal in class AbstractBlock

forward

public NDList forward(ParameterStore ps,
                      NDList inputs,
                      boolean training)

Applies the operating function of the block once. This method should be called only on blocks that are initialized.

Parameters:

ps - the parameter store

inputs - the input NDList

training - true for a training forward pass

Returns:

the output of the forward pass

forward

public NDList forward(ParameterStore ps,
                      NDArray tokenIds,
                      NDArray typeIds,
                      NDArray masks,
                      boolean training)

Embeds the input, runs it through a transformer encoder stack and outputs the resulting embeddings and the pooled classifier.

Parameters:

ps - the parameter store

tokenIds - Ids for the tokens (word pieces, BPE pieces, words etc.) (B, S)

typeIds - Ids for the token type, during pretraining this is used to encode the first/second sentence (B, S)

masks - Masks for the input, used to mask out tokens in input that is shorter than the max number of input tokens (B, S)

training - true=apply dropout etc.

Returns:

token embeddings (B, S, E), pooled classifier (B, E)

builder

public static BertBlock.Builder builder()

Returns a new BertBlock builder.

Returns:

a new BertBlock builder.