InfiniteVocabEmbedding

class torch_brain.nn.InfiniteVocabEmbedding(embedding_dim, init_scale=0.02)

Bases: torch.nn.modules.module.Module

Embedding layer with a vocabulary that can be extended. Vocabulary is saved along with the model, and is reloaded when the state_dict is loaded. This is useful when the vocabulary is dynamically generated, e.g. from a dataset. For this reason this class also plays the role of the tokenizer.

This layer is initially lazy, i.e. it does not have a weight matrix. The weight matrix is initialized when:

• The vocabulary is initialized via initialize_vocab().

• or the model is loaded from a checkpoint that contains the vocabulary.

If the vocabulary is initialized before load_state_dict() is called, an error will be raised if the vocabulary in the checkpoint does not match the vocabulary in the model. The order of the words in the vocabulary does not matter, as long as the words are the same.

If you would like to create a new variant of an existing InfiniteVocabEmbedding (that you loaded from a checkpoint), you can use:

• extend_vocab() to add new words to the vocabulary. The embeddings for the new words will be initialized randomly.

• subset_vocab() to select a subset of the vocabulary. The embeddings for the selected words will be copied from the original embeddings, and the ids for the selected words will change and tokenizer() will be updated accordingly.

This module also plays the role of the tokenizer, which is accessible via tokenizer(), and is a Callable.

Warning

If you are only interested in loading a subset of words from a checkpoint, do not call initialize_vocab(), first load the checkpoint then use subset_vocab().

Parameters:

• embedding_dim (int) – Embedding dimension.

• init_scale (float) – The standard deviation of the normal distribution used to initialize the embedding matrix. Default is 0.02.

initialize_vocab(vocab)

Initialize the vocabulary with a list of words. This method should be called only once, and before the model is trained. If you would like to add new words to the vocabulary, use extend_vocab() instead.

Note

A special word “NA” will always be in the vocabulary, and is assigned the index 0. 0 is used for padding.

Parameters:

vocab (List[str]) – A list of words to initialize the vocabulary.

Example

>>> from torch_brain.nn import InfiniteVocabEmbedding

>>> embedding = InfiniteVocabEmbedding(64)

>>> vocab = ["apple", "banana", "cherry"]
>>> embedding.initialize_vocab(vocab)

>>> embedding.vocab
OrderedDict([('NA', 0), ('apple', 1), ('banana', 2), ('cherry', 3)])

>>> embedding.weight.shape
torch.Size([4, 64])

extend_vocab(vocab, exist_ok=False)

Extend the vocabulary with a list of words. If a word already exists in the vocabulary, an error will be raised. The embeddings for the new words will be initialized randomly, and new ids will be assigned to the new words.

Parameters:

• vocab (List[str]) – A list of words to add to the vocabulary.

• exist_ok (bool) – If True, the method will not raise an error if the new words already exist in the vocabulary and will skip them. Default is False.

Example

>>> from torch_brain.nn import InfiniteVocabEmbedding

>>> embedding = InfiniteVocabEmbedding(64)

>>> vocab = ["apple", "banana", "cherry"]
>>> embedding.initialize_vocab(vocab)
>>> embedding
InfiniteVocabEmbedding(embedding_dim=64, num_embeddings=4)

>>> new_words = ["date", "elderberry", "fig"]
>>> embedding.extend_vocab(new_words)
InfiniteVocabEmbedding(embedding_dim=64, num_embeddings=7)

>>> embedding.vocab
OrderedDict([('NA', 0), ('apple', 1), ('banana', 2), ('cherry', 3), ('date', 4), ('elderberry', 5), ('fig', 6)])

>>> embedding.weight.shape
torch.Size([7, 64])

subset_vocab(vocab, inplace=True)

Select a subset of the vocabulary. The embeddings for the selected words will be copied from the original embeddings, and the ids for the selected words will be updated accordingly.

An error will be raised if one of the words does not exist in the vocabulary.

Parameters:

• vocab (List[str]) – A list of words to select from the vocabulary.

• inplace (bool) – If True, the method will modify the vocabulary and the weight matrix in place. If False, a new InfiniteVocabEmbedding will be returned with the selected words. Default is True.

Example

>>> from torch_brain.nn import InfiniteVocabEmbedding

>>> embedding = InfiniteVocabEmbedding(64)

>>> vocab = ["apple", "banana", "cherry"]
>>> embedding.initialize_vocab(vocab)
>>> embedding
InfiniteVocabEmbedding(embedding_dim=64, num_embeddings=4)

>>> selected_words = ["banana", "cherry"]
>>> embedding.subset_vocab(selected_words)
InfiniteVocabEmbedding(embedding_dim=64, num_embeddings=3)

>>> embedding.vocab
OrderedDict([('NA', 0), ('banana', 1), ('cherry', 2)])

>>> embedding.weight.shape
torch.Size([3, 64])

tokenizer(words)

Convert a word or a list of words to their token indices.

Parameters:

words (Union[str, List[str]]) – A word or a list of words.

Returns:

A token index or a list of token indices.

Return type:

Union[int, List[int]]

Example

>>> from torch_brain.nn import InfiniteVocabEmbedding

>>> embedding = InfiniteVocabEmbedding(64)

>>> vocab = ["apple", "banana", "cherry"]
>>> embedding.initialize_vocab(vocab)

>>> embedding.tokenizer("banana")
2

>>> embedding.tokenizer(["apple", "cherry", "apple"])
[1, 3, 1]

detokenizer(index)

Convert a token index to a word.

Parameters:

index (int) – A token index.

Returns:

A word.

Return type:

str

Example

>>> from torch_brain.nn import InfiniteVocabEmbedding

>>> embedding = InfiniteVocabEmbedding(64)

>>> vocab = ["apple", "banana", "cherry"]
>>> embedding.initialize_vocab(vocab)

>>> embedding.detokenizer(2)
'banana'

is_lazy()

Returns True if the module is not initialized.

Example

>>> from torch_brain.nn import InfiniteVocabEmbedding

>>> embedding = InfiniteVocabEmbedding(64)

>>> embedding.is_lazy()
True

>>> vocab = ["apple", "banana", "cherry"]
>>> embedding.initialize_vocab(vocab)

>>> embedding.is_lazy()
False

reset_parameters()

Resets all learnable parameters of the module, but will not reset the vocabulary.

forward(input)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

extra_repr()

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str