93 lines
2.9 KiB
Python
93 lines
2.9 KiB
Python
from functools import reduce
|
|
from pathlib import Path
|
|
import pytest
|
|
import Project_Model.Libs.BPE as BPE
|
|
import Project_Model.Libs.Transformer as Transformer
|
|
|
|
VOCABULARY_PATH = Path("Assets/Model/toy_10/toy_dictionary.json")
|
|
VOCABULARY = BPE.load_nanos_vocabulary(VOCABULARY_PATH)
|
|
SPECIAL_LIST = BPE.default_special_tokens()
|
|
|
|
|
|
class TestSpannedMasker:
|
|
|
|
def test_spanned_masking(self):
|
|
|
|
CORPUS_PATH = Path("Project_Model/Tests/spanner_file/mask.txt")
|
|
TEXT = CORPUS_PATH.read_text("utf-8")
|
|
CORRUPTION_PERCENTAGE = 0.15
|
|
TOLERANCE = 0.15
|
|
|
|
TOKENIZER = BPE.TokeNanoCore(VOCABULARY, SPECIAL_LIST)
|
|
VOCABULARY_SIZE = TOKENIZER.vocabulary_size
|
|
|
|
|
|
|
|
TOKENS = TOKENIZER.encode(TEXT)
|
|
|
|
LEGAL_TOKENS: set[int] = set(TOKENIZER.encode("<SUBJ><OBJ><PRED>"))
|
|
|
|
SPECIAL_TOKENS: set[int] = set(TOKENIZER.encode("".join(SPECIAL_LIST)))
|
|
|
|
ILLEGAL_TOKENS: set[int] = SPECIAL_TOKENS.difference(LEGAL_TOKENS)
|
|
|
|
MASKER = Transformer.SpannedMasker(VOCABULARY_SIZE,ILLEGAL_TOKENS,CORRUPTION_PERCENTAGE, 3)
|
|
|
|
SPECIAL_FORMATTER = TOKENIZER.encode("*<SOT>")[0]
|
|
END_FORMATTER = TOKENIZER.encode("<EOT>")[0]
|
|
|
|
OUTPUT, TARGET = MASKER.mask_sequence(TOKENS)
|
|
|
|
UNCORRUPTED_TOKENS = list(
|
|
filter(lambda token: token <= VOCABULARY_SIZE, OUTPUT)
|
|
)
|
|
CORRUPTED_TOKENS = list(filter(lambda token: token <= VOCABULARY_SIZE, TARGET))
|
|
|
|
TARGET.append(END_FORMATTER)
|
|
|
|
OUTPUT = list(
|
|
map(
|
|
lambda token: SPECIAL_FORMATTER if token > VOCABULARY_SIZE else token,
|
|
OUTPUT,
|
|
)
|
|
)
|
|
TARGET = list(
|
|
map(
|
|
lambda token: SPECIAL_FORMATTER if token > VOCABULARY_SIZE else token,
|
|
TARGET,
|
|
)
|
|
)
|
|
|
|
OUT_TEXT = TOKENIZER.decode(OUTPUT)
|
|
TAR_TEXT = TOKENIZER.decode(TARGET)
|
|
|
|
ACTUAL_CORRUPTION_PERCENTAGE = len(CORRUPTED_TOKENS) / len(TOKENS)
|
|
|
|
print(f"Original text:\n\n{TEXT}")
|
|
print(f"Inputs:\n\n{OUT_TEXT}")
|
|
print(f"Targets:\n\n{TAR_TEXT}")
|
|
print(f"Target Tokens:\n\n{OUTPUT}")
|
|
|
|
print(
|
|
"\n".join(
|
|
[
|
|
f"======================",
|
|
f"Original length: {len(TOKENS)}",
|
|
f"Uncorrupted Chars: {len(UNCORRUPTED_TOKENS)}",
|
|
f"Corrupted Chars: {len(CORRUPTED_TOKENS)}",
|
|
f"Percentage_corruption: {(len(CORRUPTED_TOKENS)/len(TOKENS))*100}%",
|
|
f"======================",
|
|
]
|
|
)
|
|
)
|
|
|
|
for token in TARGET[:len(TARGET) - 1]:
|
|
assert token not in ILLEGAL_TOKENS
|
|
|
|
assert ACTUAL_CORRUPTION_PERCENTAGE > CORRUPTION_PERCENTAGE - TOLERANCE
|
|
assert ACTUAL_CORRUPTION_PERCENTAGE < CORRUPTION_PERCENTAGE + TOLERANCE
|
|
|
|
|
|
if __name__ == "__main__":
|
|
TestSpannedMasker().test_spanned_masking()
|