78 lines
1.9 KiB
Python
78 lines
1.9 KiB
Python
from pathlib import Path
|
|
from Project_Model.Libs.BPE.Enums import TokenType
|
|
import Project_Model.Libs.BPE as BPE
|
|
|
|
import re
|
|
|
|
CACHE_DIR_PATH = Path("Project_Model/Tests/trainer_files/cache/pool-cache.json")
|
|
|
|
class TestTrainBPE:
|
|
|
|
def test_bpe_train_encoding_simple(self):
|
|
|
|
TRAINER = BPE.NanoSocraTrainerPool(
|
|
int(32E3),
|
|
["<SOT>", "<EOT>"]
|
|
)
|
|
|
|
TEXT = "abababab"
|
|
TEXT_PATH = Path("Project_Model/Tests/trainer_files/train_simple.txt")
|
|
|
|
EXPECTED = [258]
|
|
|
|
# ab = 256
|
|
# 256, 256 = 257
|
|
# 257, 257 = 258
|
|
|
|
BPE_ENCODER = TRAINER.trainBPE(
|
|
TEXT_PATH,
|
|
CACHE_DIR_PATH
|
|
)
|
|
|
|
ENCODED = BPE_ENCODER.encode(TEXT)
|
|
|
|
assert len(ENCODED) == len(EXPECTED)
|
|
|
|
for encoded, expected in zip(ENCODED, EXPECTED):
|
|
assert encoded == expected
|
|
|
|
|
|
def test_bpe_train_encoding_and_decoding(self):
|
|
|
|
SPECIAL_LIST = ["<ABS>", "<SOTL>"]
|
|
TRAINER = BPE.NanoSocraTrainerPool(
|
|
int(32E3),
|
|
SPECIAL_LIST
|
|
)
|
|
|
|
TEXT_PATH = Path("Project_Model/Tests/trainer_files/train_encode_decode.txt")
|
|
FILE = open(TEXT_PATH)
|
|
TEXT = FILE.read()
|
|
FILE.close()
|
|
|
|
EXPECTED = TEXT
|
|
|
|
# ab = 256
|
|
# 256, 256 = 257
|
|
# 257, 257 = 258
|
|
|
|
BPE_ENCODER = TRAINER.trainBPE(
|
|
TEXT_PATH,
|
|
CACHE_DIR_PATH
|
|
)
|
|
VOCABULARY = BPE_ENCODER.vocabulary
|
|
TOKENANO = BPE.TokeNanoCore(VOCABULARY,SPECIAL_LIST)
|
|
|
|
ENCODED = TOKENANO.encode(TEXT)
|
|
DECODED = TOKENANO.decode(ENCODED)
|
|
|
|
assert len(DECODED) == len(EXPECTED)
|
|
|
|
for decoded, expected in zip(DECODED, EXPECTED):
|
|
assert decoded == expected
|
|
|
|
# Useful to debug weird cases
|
|
if __name__ == "__main__":
|
|
# TestTrainBPE().test_bpe_train_encoding_simple()
|
|
TestTrainBPE().test_bpe_train_encoding_and_decoding()
|