Add test for lengths2mask
@@ -4,7 +4,7 @@ from numpy.testing import assert_equal
from spacy.tokens import Doc
from spacy.vocab import Vocab
from spacy_pytorch_transformers.pipeline.wordpiecer import align_word_pieces
-from spacy_pytorch_transformers.util import pad_batch, batch_by_length
+from spacy_pytorch_transformers.util import pad_batch, batch_by_length, lengths2mask
@pytest.mark.parametrize(
@@ -87,6 +87,21 @@ def test_pad_batch_2d(lengths, width, expected_shape):
assert padded.shape == expected_shape
+@pytest.mark.parametrize(
+ "lengths",
+ [
+ [1],
+ [1, 2],
+ [5, 3, 1],
+ ],
+)
+def test_lengths2mask(lengths):
+ mask = lengths2mask(lengths)
+ assert mask.ndim == 1
+ assert mask.sum() == sum(lengths)
+ assert mask.size == max(lengths) * len(lengths)
+
+
@pytest.mark.parametrize(
"wp_tokens,span,expected_start",
[
