ai.djl.nn.transformer

Class BertPretrainingBlock

java.lang.Object

ai.djl.nn.AbstractBlock

ai.djl.nn.transformer.BertPretrainingBlock

All Implemented Interfaces:

Block

public class BertPretrainingBlock
extends AbstractBlock

Creates a block that performs all bert pretraining tasks (next sentence and masked language).

Field Summary

Fields inherited from class ai.djl.nn.AbstractBlock

children, inputNames, inputShapes, parameters, parameterShapeCallbacks, version

Constructor Summary

Constructors

Constructor and Description
BertPretrainingBlock(BertBlock.Builder builder) Creates a new Bert pretraining block fitting the given bert configuration.

Method Summary

Modifier and Type	Method and Description
NDList	forward(ParameterStore ps, NDArray tokenIds, NDArray typeIds, NDArray sequenceMasks, NDArray maskedIndices, boolean training) Applies one bert pretraining step.
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)
Shape[]	getOutputShapes(NDManager manager, Shape[] inputShapes) Returns the expected output shapes of the block for the specified input shapes.
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

Constructor Detail

BertPretrainingBlock

public BertPretrainingBlock(BertBlock.Builder builder)

Creates a new Bert pretraining block fitting the given bert configuration.

Parameters:

builder - a builder with a bert configuration

Method Detail

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

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 sequenceMasks,
                      NDArray maskedIndices,
                      boolean training)

Applies one bert pretraining step.

Parameters:

ps - the parameter store

tokenIds - int, (B, S)

typeIds - int, (B, S)

sequenceMasks - int, (B, S)

maskedIndices - int, (B, I)

training - true=apply dropout etc.

Returns:

next sentence probabilities (B, 2), masked token probabilities (B, I, D),

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