public final class BertBlock extends AbstractBlock
This closely follows the original Devlin et. al. paper and its reference implementation.
Modifier and Type | Class and Description |
---|---|
static class |
BertBlock.Builder
|
children, inputNames, inputShapes, parameters, version
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.
|
protected NDList |
forwardInternal(ParameterStore ps,
NDList inputs,
boolean training,
ai.djl.util.PairList<java.lang.String,java.lang.Object> params)
A helper for
Block.forward(ParameterStore, NDList, boolean, PairList) after
initialization. |
int |
getEmbeddingSize()
Returns the embedding size used for tokens.
|
Shape[] |
getOutputShapes(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.
|
addChildBlock, addParameter, beforeInitialize, cast, clear, describeInput, forward, forward, forwardInternal, getChildren, getDirectParameters, getParameters, initialize, isInitialized, loadMetadata, loadParameters, prepare, readInputShapes, saveInputShapes, saveMetadata, saveParameters, setInitializer, setInitializer, setInitializer, toString
clone, equals, finalize, getClass, hashCode, notify, notifyAll, wait, wait, wait
forward, validateLayout
public IdEmbedding getTokenEmbedding()
public int getEmbeddingSize()
public int getTokenDictionarySize()
public int getTypeDictionarySize()
public Shape[] getOutputShapes(Shape[] inputShapes)
inputShapes
- the shapes of the inputspublic void initializeChildBlocks(NDManager manager, DataType dataType, Shape... inputShapes)
initializeChildBlocks
in class AbstractBlock
manager
- the manager to use for initializationdataType
- the requested data typeinputShapes
- the expected input shapes for this blockpublic static NDArray createAttentionMaskFromInputMask(NDArray ids, NDArray mask)
ids
- 2D Tensor of shape (B, F)mask
- 2D Tensor of shape (B, T)protected NDList forwardInternal(ParameterStore ps, NDList inputs, boolean training, ai.djl.util.PairList<java.lang.String,java.lang.Object> params)
Block.forward(ParameterStore, NDList, boolean, PairList)
after
initialization.forwardInternal
in class AbstractBlock
ps
- the parameter storeinputs
- the input NDListtraining
- true for a training forward passparams
- optional parameterspublic static BertBlock.Builder builder()