get_position_ids_from_padding_mask

torchtune.generation.get_position_ids_from_padding_mask(padding_mask: Tensor)

Calculates position ids given a padding mask which right-shifts position ids to start from the first valid token.

Parameters:

padding_mask (torch.Tensor) – Boolean tensor where False indicates the corresponding token in the sequence is a padding token and should be masked out in attention. Shape [bsz, seq_len]

Returns:

position ids which are appropriately shifted according to any padding values.

Return type:

torch.Tensor

Example

>>> padding_mask = torch.tensor([False, False, False, True, True, True, True, True])
>>> get_position_ids_from_padding_mask(padding_mask)
torch.Tensor([0, 0, 0, 0, 1, 2, 3, 4])