Add a `max_words` argument to `build_vocab_from_iterator`

[xhluca]

🚀 Feature

Link to the docs

I believe it would be beneficial to limit the number of words you want in your vocabulary with an argument like max_words, e.g.:

vocab = build_vocab_from_iterator(yield_tokens_batch(file_path), specials=["<unk>"], max_words=50000)

Motivation

This allows a controllable-sized nn.Embedding, with rare words being mapped to <unk>. Otherwise, it would not be practical to use build_vocab_from_iterator for larger datasets.

Alternatives

Keras and Huggingface's tokenizers would be viable alternatives, but do not nicely integrate with the torchtext ecosystem.

[parmeet]

Thank you @xhlulu for interesting proposal. I don't see any downsides for adding this argument. Internally, we are sorting tokens by frequency and I think it would be trivial to add the support for max_words by discarding the lowest frequency tokens. I wonder if there are any other criteria that users might want to follow instead of frequency?

[erip]

I'll add that normally your "upstream" tokenization strategy will handle this for you. For instance, you can specify subword vocab size when training unigram LM as in sentencepiece or BPE-based models. Building your torchtext vocabulary from subword tokenized segments means you already get this "max_tokens" implicitly.

In python-y pseudocode, you could do something like...

# subword_model is a pretrained subword model w/ pre-specified vocab size max_tokens
def yield_tokens(file_path, subword_model):
    for line in get_lines_from_file(file_path):
        yield from subword_model.tokenize(line)

vocab = build_vocab_from_iterator(yield_tokens(file_path, subword_model), specials=[...])
assert len(vocab) <=  max_tokens
embedding = nn.Embedding(len(vocab), embedding_dim, padding_index=vocab["<pad>"])

[parmeet]

I'll add that normally your "upstream" tokenization strategy will handle this for you.

I wonder if this is valid for all types of tokenizers? For instance white-space or other rule/regex based tokenizers may not provide such constraints out-of-the-box, right?

[erip]

It's certainly not valid in all cases and your point is well taken. I'd at the very least recommend sticking with "max tokens" or "max vocab size" since words are a bit sticky. :-)

[abhinavarora]

This does seem like a valid request as this limit from upstream tokenizers is something that is not applicable to regex, whitespace or from a pre-existing vocabulary. Agree with @erip on the naming being something like max_vocab_size

[abhinavarora]

Thank you @xhlulu for interesting proposal. I don't see any downsides for adding this argument. Internally, we are sorting tokens by frequency and I think it would be trivial to add the support for max_words by discarding the lowest frequency tokens. I wonder if there are any other criteria that users might want to follow instead of frequency?

Having the frequency criteris seems good enough for now :-)