ai.djl.modality.nlp.embedding

Class TrainableWordEmbedding

java.lang.Object

ai.djl.nn.AbstractBlock

ai.djl.nn.ParameterBlock

ai.djl.nn.core.Embedding<java.lang.String>

ai.djl.modality.nlp.embedding.TrainableWordEmbedding

All Implemented Interfaces:

WordEmbedding, Block, AbstractEmbedding<java.lang.String>, AbstractIndexedEmbedding<java.lang.String>

public class TrainableWordEmbedding
extends Embedding<java.lang.String>
implements WordEmbedding

TrainableWordEmbedding is an implementation of WordEmbedding and Embedding based on a SimpleVocabulary. This WordEmbedding is ideal when there are no pre-trained embeddings available.

Nested Class Summary

Nested Classes

Modifier and Type	Class and Description
static class	TrainableWordEmbedding.Builder A builder for a TrainableWordEmbedding.

Nested classes/interfaces inherited from class ai.djl.nn.core.Embedding

Embedding.BaseBuilder<T,B extends Embedding.BaseBuilder<T,B>>, Embedding.DefaultEmbedding, Embedding.DefaultItem

Field Summary

Fields inherited from class ai.djl.nn.core.Embedding

dataType, embedder, embedding, embeddingSize, fallthroughEmbedding, numItems, sparseGrad, unembedder

Fields inherited from class ai.djl.nn.AbstractBlock

inputNames, inputShapes

Constructor Summary

Constructors

Constructor and Description
TrainableWordEmbedding(NDArray embedding, java.util.List<java.lang.String> items) Constructs a pretrained embedding.
TrainableWordEmbedding(NDArray embedding, java.util.List<java.lang.String> items, boolean sparseGrad) Constructs a pretrained embedding.
TrainableWordEmbedding(SimpleVocabulary simpleVocabulary, int embeddingSize) Constructs a new instance of TrainableWordEmbedding from a SimpleVocabulary and a given embedding size.
TrainableWordEmbedding(TrainableWordEmbedding.Builder builder) Constructs a new instance of TrainableWordEmbedding from the TrainableWordEmbedding.Builder.

Method Summary

Modifier and Type	Method and Description
static TrainableWordEmbedding.Builder	builder() Creates a builder to build an Embedding.
java.lang.String	decode(byte[] byteArray) Decodes the given byte array into an object of input parameter type.
NDArray	embedWord(NDManager manager, int index) Embeds the word after preprocessed using WordEmbedding.preprocessWordToEmbed(String).
byte[]	encode(java.lang.String input) Encodes an object of input type into a byte array.
int	preprocessWordToEmbed(java.lang.String word) Pre-processes the word to embed into an array to pass into the model.
java.lang.String	unembedWord(NDArray word) Returns the closest matching word for the given index.
boolean	vocabularyContains(java.lang.String word) Returns whether an embedding exists for a word.

Methods inherited from class ai.djl.nn.core.Embedding

embed, embed, forward, getDirectParameters, getOutputShapes, getParameterShape, hasItem, loadParameters, saveParameters, unembed

Methods inherited from class ai.djl.nn.ParameterBlock

getChildren, initialize, toString

Methods inherited from class ai.djl.nn.AbstractBlock

beforeInitialize, cast, clear, describeInput, getParameters, isInitialized, readInputShapes, saveInputShapes, setInitializer, setInitializer

Methods inherited from class java.lang.Object

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

Methods inherited from interface ai.djl.modality.nlp.embedding.WordEmbedding

embedWord

Methods inherited from interface ai.djl.nn.Block

forward, validateLayout

Constructor Detail

TrainableWordEmbedding

public TrainableWordEmbedding(TrainableWordEmbedding.Builder builder)

Constructs a new instance of TrainableWordEmbedding from the TrainableWordEmbedding.Builder.

Parameters:

builder - the TrainableWordEmbedding.Builder

TrainableWordEmbedding

public TrainableWordEmbedding(SimpleVocabulary simpleVocabulary,
                              int embeddingSize)

Constructs a new instance of TrainableWordEmbedding from a SimpleVocabulary and a given embedding size.

Parameters:

simpleVocabulary - a SimpleVocabulary to get tokens from

embeddingSize - the required embedding size

TrainableWordEmbedding

public TrainableWordEmbedding(NDArray embedding,
                              java.util.List<java.lang.String> items)

Constructs a pretrained embedding.

Parameters:

embedding - the embedding array

items - the items in the embedding (in matching order to the embedding array)

TrainableWordEmbedding

public TrainableWordEmbedding(NDArray embedding,
                              java.util.List<java.lang.String> items,
                              boolean sparseGrad)

Constructs a pretrained embedding.

Parameters:

embedding - the embedding array

items - the items in the embedding (in matching order to the embedding array)

sparseGrad - whether to compute row sparse gradient in the backward calculation

Method Detail

vocabularyContains

public boolean vocabularyContains(java.lang.String word)

Returns whether an embedding exists for a word.

Specified by:

vocabularyContains in interface WordEmbedding

Parameters:

word - the word to check

Returns:

true if an embedding exists

preprocessWordToEmbed

public int preprocessWordToEmbed(java.lang.String word)

Pre-processes the word to embed into an array to pass into the model.

Make sure to call WordEmbedding.embedWord(NDManager, int) after this.

Specified by:

preprocessWordToEmbed in interface WordEmbedding

Parameters:

word - the word to embed

Returns:

the word that is ready to embed

embedWord

public NDArray embedWord(NDManager manager,
                         int index)

Embeds the word after preprocessed using WordEmbedding.preprocessWordToEmbed(String).

Specified by:

embedWord in interface WordEmbedding

Parameters:

manager - the manager for the embedding array

index - the index of the word to embed

Returns:

the embedded word

unembedWord

public java.lang.String unembedWord(NDArray word)

Returns the closest matching word for the given index.

Specified by:

unembedWord in interface WordEmbedding

Parameters:

word - the word embedding to find the matching string word for.

Returns:

a word similar to the passed in embedding

encode

public byte[] encode(java.lang.String input)

Encodes an object of input type into a byte array. This is used in saving and loading the Embedding objects.

Specified by:

encode in interface AbstractIndexedEmbedding<java.lang.String>

Parameters:

input - the input object to be encoded

Returns:

the encoded byte array.

decode

public java.lang.String decode(byte[] byteArray)

Decodes the given byte array into an object of input parameter type.

Specified by:

decode in interface AbstractIndexedEmbedding<java.lang.String>

Parameters:

byteArray - the byte array to be decoded

Returns:

the decode object of input parameter type

builder

public static TrainableWordEmbedding.Builder builder()

Creates a builder to build an Embedding.

Returns:

a new builder