Feature Request: Masking

(more requests from the trenches, as always love the work!)

I would love if torchtext handled masking for me. It’s a small think, but I don’t like that I have to pass around pad_index in my code. It breaks the abstraction.

• Idea: batch object should have masks. Could we add something like this in addition to lengths or maybe just along with them?

def sequence_mask(lengths, max_len=None):
    """
    Creates a boolean mask from sequence lengths.
    """
    batch_size = lengths.numel()
    max_len = max_len or lengths.max()
    return (torch.arange(0, max_len)
            .type_as(lengths)
            .repeat(batch_size, 1)
            .lt(lengths.unsqueeze(1)))

• Another really useful mask is to blackout future words along with padding. We do something like this (although this code looks too complicated.

    def _get_attn_subsequent_mask(self, size):
        ''' Get an attention mask to avoid using the subsequent info.'''
        attn_shape = (1, size, size)
        subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
        subsequent_mask = torch.from_numpy(subsequent_mask)
        return subsequent_mask

[jekbradbury]

This is coming soon, in a separate package called Matchbox that we are in the process of open-sourcing! (It does a little more than you’re looking for: it introduces a batch type that carries an associated mask, but also overloads all PyTorch operations on this type so that the mask is propagated correctly.)

[jekbradbury]

The main reason Matchbox didn’t get mature enough was that it became clear that, while doing it in C++/TorchScript might have reasonable performance, doing it in Python probably never would. The main reason BatchTensor didn’t get mature enough (as far as I can tell) was that it was pretty difficult to work on the JIT at a time when everything was changing constantly, and because the JIT IR wasn’t expressive enough at the time to do it cleanly. I also stopped working on these things when I moved to Google, but I don’t think that’s as big a factor as the above two. But Soumith mentioned a new BatchTensor project in Slack the other day, so I’m excited about that!