Fix test with special tokens
import pytest
from spacy_pytorch_transformers import PyTT_WordPiecer
+from spacy_pytorch_transformers.util import is_special_token
from spacy.vocab import Vocab
from spacy.tokens import Doc
@@ -15,4 +16,6 @@ def test_wordpiecer(wp):
doc[0].is_sent_start = True
doc[1].is_sent_start = False
doc = wp(doc)
- assert wp.model.clean_wp_tokens(doc._.pytt_word_pieces_) == words
+ cleaned_words = wp.model.clean_wp_tokens(doc._.pytt_word_pieces_)
+ cleaned_words = [w for w in cleaned_words if not is_special_token(w)]
+ assert cleaned_words == words
