Skip to content

Instantly share code, notes, and snippets.

@binshengliu

binshengliu/len_to_mask.py

Last active August 13, 2020 01:43
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save binshengliu/23e516a2ac596b37a61130c66c4a4a54 to your computer and use it in GitHub Desktop.
Save binshengliu/23e516a2ac596b37a61130c66c4a4a54 to your computer and use it in GitHub Desktop.
Download ZIP
Convert different lengths into mask tensor
Raw
len_to_mask.py
# lens = [3, 5, 4]
# What we want:
# mask = [[1, 1, 1, 0, 0],
# [1, 1, 1, 1, 1],
# [1, 1, 1, 1, 0]]
# https://stackoverflow.com/questions/53403306/how-to-batch-convert-sentence-lengths-to-masks-in-pytorch
def len_to_mask(lens: np.ndarray, seq_len: Optional[int] = None) -> np.ndarray:
if seq_len is None:
seq_len = max(lens)
return np.arange(seq_len)[None, :] < lens[:, None]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment