39 lines
1.3 KiB
Python
39 lines
1.3 KiB
Python
import pytest
|
||
from serializer import Serializer
|
||
from vocab import Vocab
|
||
|
||
|
||
def test_vocab():
|
||
vocab = Vocab('test')
|
||
sent = ' 我是中国人，我爱中国。 I\'m Chinese, I love China'
|
||
|
||
serializer = Serializer(do_lower_case=True)
|
||
tokens = serializer.serialize(sent)
|
||
assert tokens == [
|
||
'我', '是', '中', '国', '人', '，', '我', '爱', '中', '国', '。', 'i', "'", 'm', 'chinese', ',', 'i', 'love', 'china'
|
||
]
|
||
|
||
vocab.add_words(tokens)
|
||
unk_str = '[UNK]'
|
||
unk_idx = vocab.word2idx[unk_str]
|
||
|
||
assert vocab.count == 22
|
||
assert len(vocab.word2idx) == len(vocab.idx2word) == len(vocab.word2idx) == 22
|
||
|
||
vocab.trim(2, verbose=False)
|
||
|
||
assert vocab.count == 11
|
||
assert len(vocab.word2idx) == len(vocab.idx2word) == len(vocab.word2idx) == 11
|
||
|
||
token2idx = [vocab.word2idx.get(i, unk_idx) for i in tokens]
|
||
assert len(tokens) == len(token2idx)
|
||
assert token2idx == [7, 1, 8, 9, 1, 1, 7, 1, 8, 9, 1, 10, 1, 1, 1, 1, 10, 1, 1]
|
||
|
||
idx2tokens = [vocab.idx2word.get(i, unk_str) for i in token2idx]
|
||
assert len(idx2tokens) == len(token2idx)
|
||
assert ' '.join(idx2tokens) == '我 [UNK] 中 国 [UNK] [UNK] 我 [UNK] 中 国 [UNK] i [UNK] [UNK] [UNK] [UNK] i [UNK] [UNK]'
|
||
|
||
|
||
if __name__ == '__main__':
|
||
pytest.main()
|