Fix import and formatting
@@ -3,8 +3,8 @@ import numpy
from numpy.testing import assert_equal
from spacy.tokens import Doc
from spacy.vocab import Vocab
-from spacy_pytorch_transformers.util import align_word_pieces, pad_batch
-from spacy_pytorch_transformers.util import batch_by_length
+from spacy_pytorch_transformers.wordpiecer import align_word_pieces
+from spacy_pytorch_transformers.util import pad_batch, batch_by_length
@pytest.mark.parametrize(
@@ -18,9 +18,12 @@ from spacy_pytorch_transformers.util import batch_by_length
(["abcd"], ["ab", "cd"], [[0, 1]]),
(["d", "e", "f"], ["[CLS]", "d", "e", "f"], [[1], [2], [3]]),
([], [], []),
- (["it", "realllllllly", "sucks", "."], ["it", "real", "ll", "ll", "ll", "ly", "suck", "s", "."],
- [[0], [1, 2, 3, 4, 5], [6, 7], [8]]),
- (["Well", ",", "i"], ['Well', ',', '', 'i'], [[0], [1, 2], [2, 3]])
+ (
+ ["it", "realllllllly", "sucks", "."],
+ ["it", "real", "ll", "ll", "ll", "ly", "suck", "s", "."],
+ [[0], [1, 2, 3, 4, 5], [6, 7], [8]],
+ ),
+ (["Well", ",", "i"], ["Well", ",", "", "i"], [[0], [1, 2], [2, 3]]),
],
)
def test_align_word_pieces(spacy_tokens, wp_tokens, expected_alignment):
