|
import os
|
|
import pytest
|
|
from bpetokenizer import BPETokenizer, Tokenizer
|
|
|
|
@pytest.fixture
|
|
def tokenizer():
|
|
return Tokenizer()
|
|
|
|
@pytest.fixture
|
|
def bpe_tokenizer():
|
|
return BPETokenizer()
|
|
|
|
|
|
def test_train():
|
|
"""Test the training of the tokenizer."""
|
|
text = "aaabdaaabac"
|
|
tokenizer = Tokenizer()
|
|
tokenizer.train(text, 259, verbose=False)
|
|
assert len(tokenizer.vocab) == 259
|
|
assert len(tokenizer.merges) == 3
|
|
assert tokenizer.decode(tokenizer.encode(text)) == "aaabdaaabac"
|
|
|
|
|
|
def test_encode():
|
|
"""Test the encoding of the tokenizer."""
|
|
text = "aaabdaaabac"
|
|
tokenizer = Tokenizer()
|
|
tokenizer.train(text, 259, verbose=False)
|
|
assert tokenizer.encode("aaabdaaabac") == [258, 100, 258, 97, 99]
|
|
|
|
|
|
def test_decode():
|
|
"""Test the decoding of the tokenizer."""
|
|
text = "aaabdaaabac"
|
|
tokenizer = Tokenizer()
|
|
tokenizer.train(text, 259, verbose=False)
|
|
assert tokenizer.decode([258, 100, 258, 97, 99]) == "aaabdaaabac"
|
|
|
|
|
|
def test_train_bpe():
|
|
"""Test the training of the BPE tokenizer."""
|
|
text = "aaabdaaabac"
|
|
tokenizer = BPETokenizer()
|
|
tokenizer.train(text, 256 + 3, verbose=False)
|
|
assert len(tokenizer.vocab) == 259
|
|
assert len(tokenizer.merges) == 3
|
|
assert tokenizer.decode(tokenizer.encode(text)) == "aaabdaaabac"
|
|
|
|
|
|
def test_train_bpe_w_special_tokens():
|
|
"""Test the bpetokenizer with special tokens"""
|
|
special_tokens = {
|
|
"<|endoftext|>": 1001,
|
|
"<|startoftext|>": 1002,
|
|
"[SPECIAL1]": 1003,
|
|
"[SPECIAL2]": 1004,
|
|
}
|
|
|
|
PATTERN = r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+"""
|
|
tokenizer = BPETokenizer(special_tokens=special_tokens, pattern=PATTERN)
|
|
texts = "<|startoftext|> Hello, World! This is a sample text with the special tokens [SPECIAL1] and [SPECIAL2] to test the tokenizer.<|endoftext|>"
|
|
tokenizer.train(texts, vocab_size=310, verbose=False)
|
|
|
|
assert len(tokenizer.vocab) == 310
|
|
assert len(tokenizer.merges) == 310 - 256
|
|
assert tokenizer.decode(tokenizer.encode(texts)) == texts
|
|
assert tokenizer.inverse_special_tokens == {v: k for k,v in special_tokens.items()}
|
|
assert tokenizer.special_tokens == special_tokens
|
|
assert tokenizer.pattern == PATTERN
|
|
|
