Package ai.djl.modality.nlp.embedding

Class ModelZooTextEmbedding

java.lang.Object

ai.djl.modality.nlp.embedding.ModelZooTextEmbedding

All Implemented Interfaces:

TextEmbedding, AutoCloseable

public class ModelZooTextEmbedding
extends Object
implements TextEmbedding, AutoCloseable

A WordEmbedding using a ZooModel.

Constructor Summary

Constructors

Constructor	Description
ModelZooTextEmbedding(Model model)	Constructs a ModelZooTextEmbedding.

Method Summary

Modifier and Type	Method	Description
void	close()
NDArray	embedText(NDArray indices)	Embeds the text after preprocessed using TextEmbedding.preprocessTextToEmbed(List).
long[]	preprocessTextToEmbed(List<String> tokens)	Preprocesses the text to embed into an array to pass into the model.
List<String>	unembedText(NDArray word)	Returns the closest matching text for a given embedding.

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.TextEmbedding

embedText, embedText

Constructor Details

ModelZooTextEmbedding

public ModelZooTextEmbedding(Model model)

Constructs a ModelZooTextEmbedding.

Parameters:

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

Method Details

preprocessTextToEmbed

public long[] preprocessTextToEmbed(List<String> tokens)

Preprocesses the text to embed into an array to pass into the model.

Make sure to call TextEmbedding.embedText(NDManager, long[]) after this.

Specified by:

preprocessTextToEmbed in interface TextEmbedding

Parameters:

tokens - the text to embed

Returns:

the indices of text that is ready to embed

embedText

public NDArray embedText(NDArray indices)
                  throws EmbeddingException

Embeds the text after preprocessed using TextEmbedding.preprocessTextToEmbed(List).

Specified by:

embedText in interface TextEmbedding

Parameters:

indices - the indices of text to embed

Returns:

the embedded text

Throws:

EmbeddingException - if there is an error while trying to embed

unembedText

public List<String> unembedText(NDArray word)

Returns the closest matching text for a given embedding.

Specified by:

unembedText in interface TextEmbedding

Parameters:

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

Returns:

text similar to the passed in embedding

close

public void close()

Specified by:

close in interface AutoCloseable