147 lines
4.6 KiB
Python
147 lines
4.6 KiB
Python
import os
|
|
import sys
|
|
import unittest
|
|
import sentencepiece as spm
|
|
import torch
|
|
|
|
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
|
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
|
|
sys.path.append(MODULE_ROOT)
|
|
|
|
from sentencepiece_wrapper.sp_load_model import SpLoadModel
|
|
from sentencepiece_wrapper.sp_encode import SpEncode
|
|
|
|
|
|
class TestSpEncode(unittest.TestCase):
|
|
def setUp(self):
|
|
self.load_node = SpLoadModel()
|
|
self.node = SpEncode()
|
|
self.sp = self.load_node.f("spiece.model")[0]
|
|
|
|
def test_1d_no_padding(self):
|
|
sentence = "Hello, world"
|
|
encoder = self.node.f(
|
|
spmodel=self.sp,
|
|
padding=False,
|
|
padding_method="longest",
|
|
padding_value=0,
|
|
truncation=False,
|
|
max_length=256
|
|
)[0]
|
|
tokens, mask = encoder(sentence=sentence)
|
|
print(tokens)
|
|
expected_tokens = torch.tensor([8774, 6, 296], dtype=torch.long)
|
|
expected_mask = torch.tensor([1, 1, 1], dtype=torch.long)
|
|
self.assertTrue(torch.equal(tokens, expected_tokens))
|
|
self.assertTrue(torch.equal(mask, expected_mask))
|
|
|
|
def test_2d_no_padding(self):
|
|
sentence = ["Hello,", "world"]
|
|
encoder = self.node.f(
|
|
spmodel=self.sp,
|
|
padding=False,
|
|
padding_method="longest",
|
|
padding_value=0,
|
|
truncation=False,
|
|
max_length=256
|
|
)[0]
|
|
tokens, masks = encoder(sentence=sentence)
|
|
|
|
expected_tokens = [
|
|
torch.tensor([8774, 6], dtype=torch.long),
|
|
torch.tensor([296], dtype=torch.long)
|
|
]
|
|
expected_masks = [
|
|
torch.tensor([1, 1], dtype=torch.long),
|
|
torch.tensor([1], dtype=torch.long)
|
|
]
|
|
for t, e in zip(tokens, expected_tokens):
|
|
self.assertTrue(torch.equal(t, e))
|
|
for m, e in zip(masks, expected_masks):
|
|
self.assertTrue(torch.equal(m, e))
|
|
|
|
def test_1d_with_padding(self):
|
|
sentence = "Hello, world"
|
|
encoder = self.node.f(
|
|
spmodel=self.sp,
|
|
padding=True,
|
|
padding_method="longest",
|
|
padding_value=0,
|
|
truncation=False,
|
|
max_length=256
|
|
)[0]
|
|
tokens, mask = encoder(sentence=sentence)
|
|
expected_tokens = torch.tensor([8774, 6, 296], dtype=torch.long)
|
|
expected_mask = torch.tensor([1, 1, 1], dtype=torch.long)
|
|
self.assertTrue(torch.equal(tokens, expected_tokens))
|
|
self.assertTrue(torch.equal(mask, expected_mask))
|
|
|
|
def test_2d_with_padding(self):
|
|
sentence = ["Hello,", "world"]
|
|
encoder = self.node.f(
|
|
spmodel=self.sp,
|
|
padding=True,
|
|
padding_method="longest",
|
|
padding_value=0,
|
|
truncation=False,
|
|
max_length=256
|
|
)[0]
|
|
tokens, mask = encoder(sentence=sentence)
|
|
|
|
expected_tokens = torch.tensor([
|
|
[8774, 6],
|
|
[296, 0]
|
|
], dtype=torch.long)
|
|
expected_mask = torch.tensor([
|
|
[1, 1],
|
|
[1, 0]
|
|
], dtype=torch.long)
|
|
self.assertTrue(torch.equal(tokens, expected_tokens))
|
|
self.assertTrue(torch.equal(mask, expected_mask))
|
|
|
|
def test_1d_truncation(self):
|
|
sentence = "Hello, world"
|
|
encoder = self.node.f(
|
|
spmodel=self.sp,
|
|
padding=False,
|
|
padding_method="longest",
|
|
padding_value=0,
|
|
truncation=True,
|
|
max_length=2
|
|
)[0]
|
|
tokens, mask = encoder(sentence=sentence)
|
|
|
|
expected_tokens = torch.tensor([8774, 6], dtype=torch.long)
|
|
expected_mask = torch.tensor([1, 1], dtype=torch.long)
|
|
self.assertTrue(torch.equal(tokens, expected_tokens))
|
|
self.assertTrue(torch.equal(mask, expected_mask))
|
|
|
|
def test_2d_truncation(self):
|
|
sentence = ["Hello,", "world"]
|
|
encoder = self.node.f(
|
|
spmodel=self.sp,
|
|
padding=False,
|
|
padding_method="longest",
|
|
padding_value=0,
|
|
truncation=True,
|
|
max_length=1
|
|
)[0]
|
|
tokens, masks = encoder(sentence=sentence)
|
|
|
|
expected_tokens = [
|
|
torch.tensor([8774], dtype=torch.long),
|
|
torch.tensor([296], dtype=torch.long)
|
|
]
|
|
expected_masks = [
|
|
torch.tensor([1], dtype=torch.long),
|
|
torch.tensor([1], dtype=torch.long)
|
|
]
|
|
for t, e in zip(tokens, expected_tokens):
|
|
self.assertTrue(torch.equal(t, e))
|
|
for m, e in zip(masks, expected_masks):
|
|
self.assertTrue(torch.equal(m, e))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|