ai.djl.modality.nlp.embedding

Class ModelZooWordEmbedding

java.lang.Object

ai.djl.modality.nlp.embedding.ModelZooWordEmbedding

All Implemented Interfaces:

WordEmbedding, java.lang.AutoCloseable

public class ModelZooWordEmbedding
extends java.lang.Object
implements WordEmbedding, java.lang.AutoCloseable

A WordEmbedding using a ZooModel.

Constructor Summary

Constructors

Constructor and Description
ModelZooWordEmbedding(Model model) Constructs a ModelZooWordEmbedding.

Method Summary

Modifier and Type	Method and Description
void	close()
NDArray	embedWord(NDManager manager, int word) Embeds the word after preprocessed using WordEmbedding.preprocessWordToEmbed(String).
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 java.lang.Object

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

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

embedWord

Constructor Detail

ModelZooWordEmbedding

public ModelZooWordEmbedding(Model model)

Constructs a ModelZooWordEmbedding.

Parameters:

model - the model for the embedding. The model's block must consist of only an Embedding<String>.

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 word)
                  throws EmbeddingException

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

Specified by:

embedWord in interface WordEmbedding

Parameters:

manager - the manager for the embedding array

word - the index of the word to embed

Returns:

the embedded word

Throws:

EmbeddingException - if there is an error while trying to embed

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

close

public void close()

Specified by:

close in interface java.lang.AutoCloseable