Test RaggedArray methods
import pytest
import numpy
-from spacy_pytorch_transformers.activations import Activations
+from spacy_pytorch_transformers.activations import Activations, RaggedArray
-@pytest.fixture
-def model(nlp):
- return nlp.get_pipe("pytt_tok2vec").model._model
-
def test_act_blank():
acts = Activations.blank()
assert acts.lh.data.size == 0
@@ -15,3 +11,35 @@ def test_act_blank():
assert acts.po.lengths == []
assert not acts.has_lh
assert not acts.has_po
+
+
+@pytest.mark.parametrize("extra_dims,lengths,pad_to,expected_shape", [
+ ((), [1], 1, (1,1)),
+ ((), [1, 2], -1, (2, 2)),
+ ((3,), [1, 2], -1, (2, 2, 3)),
+ ((3,), [1, 2, 5], -1, (3, 5, 3)),
+ ((3,), [1, 2, 5], 4, (3, 5, 3)),
+])
+def test_ragged_to_padded(extra_dims, lengths, pad_to, expected_shape):
+ arr = RaggedArray(numpy.ones((sum(lengths),) + extra_dims), lengths)
+ if pad_to < max(lengths):
+ with pytest.raises(ValueError):
+ padded = arr.to_padded(to=pad_to)
+ else:
+ padded = arr.to_padded(to=pad_to)
+ assert padded.shape == expected_shape
+
+
+@pytest.mark.parametrize("shape,lengths,expected_shape", [
+ ((1,2), [3], (3,)),
+ ((2,2), [3,2], (5,)),
+ ((2,2,4), [3,2], (5,4)),
+ ((4,2,2), [5,2,1,3], (11,2)),
+])
+def test_ragged_from_truncated(shape, lengths, expected_shape):
+ truncated = numpy.ones(shape, dtype="i")
+ for i, length in enumerate(lengths):
+ truncated[i, length:] = 0
+ ragged = RaggedArray.from_truncated(truncated, lengths)
+ assert ragged.data.shape == expected_shape
+ assert ragged.data.sum() == truncated.sum()
