Updated DistilBERT test to reflect the sequence encoding

2019-09-11 11:22:04 +02:00
parent d572d7027b
commit c4d4f3ec8c
1 changed files with 2 additions and 2 deletions
--- a/pytorch_transformers/tests/tokenization_distilbert_test.py
+++ b/pytorch_transformers/tests/tokenization_distilbert_test.py
@@ -39,8 +39,8 @@ class DistilBertTokenizationTest(BertTokenizationTest):
        encoded_sentence = tokenizer.add_special_tokens_single_sentence(text)
        encoded_pair = tokenizer.add_special_tokens_sentences_pair(text, text_2)

-        assert encoded_sentence == [101] + text + [102]
-        assert encoded_pair == [101] + text + [102] + text_2 + [102]
+        assert encoded_sentence == text
+        assert encoded_pair == text + [102] + text_2

 if __name__ == '__main__':
    unittest.main()