feat: add test for dataset

Author: j5g-moon

tests/__init__.py

@@ -1,7 +1,5 @@
 import os
 import sys
 
-sys.path.insert(
-    0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../tokenizer"))
-)
+
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

tests/test_dataset.py

@@ -0,0 +1,47 @@
+import pytest  # noqa: E902
+
+from tokenizers import SentencePieceBPETokenizer
+
+from dataset import WMT14Dataset
+
+
+@pytest.mark.parametrize("langpair", ["en-de"])
+@pytest.mark.parametrize("source_lines", [["Der Bau und die Reparatur der Autostraßen...", "die Mitteilungen sollen den geschäftlichen kommerziellen Charakter tragen."]])
+@pytest.mark.parametrize("target_lines", [["Construction and repair of highways and...", "An announcement must be commercial character."]]
+def test_len(langpair, source_lines, target_lines):
+    ds = WMT14dataset(langpair, source_lines, target_lines)
+    assert len(ds.source_lines) == len(ds)
+
+
+@pytest.mark.parametrize("langpair", ["en-de"])
+@pytest.mark.parametrize("source_lines", [["Der Bau und die Reparatur der Autostraßen...", "die Mitteilungen sollen den geschäftlichen kommerziellen Charakter tragen."]])
+@pytest.mark.parametrize("target_lines", [["Construction and repair of highways and...", "An announcement must be commercial character."]]
+def test_getitem(langpair, source_lines, target_lines):
+    ds = WMT14dataset(langpair, source_lines, target_lines)
+    source_encode_pad_test, target_encode_pad_test = ds[0]
+    assert source_encode_pad_test.size() == target_encode_pad_test.size()
+    assert source_encode_pad_test.size()[0] == ds.model_config.max_len
+
+
+@pytest.mark.parametrize("langpair", ["en-de"])
+@pytest.mark.parametrize("source_lines", [["Der Bau und die Reparatur der Autostraßen...", "die Mitteilungen sollen den geschäftlichen kommerziellen Charakter tragen."]])
+@pytest.mark.parametrize("target_lines", [["Construction and repair of highways and...", "An announcement must be commercial character."]]
+def test_encode(langpair, source_lines, target_lines):
+    ds = WMT14dataset(langpair, source_lines, target_lines)
+    source_encode_test, target_encode_test = ds._encode(source_lines[0], target_lines[0])
+    bos = ds.tokenizer.token_to_id('<bos>')
+    eos = ds.tokenizer.token_to_id('<eos>')
+    assert target_encode_test[0] == bos
+    assert target_encode_test[-1] == eos
+    assert isinstance(source_encode_test, list)
+    assert isinstance(source_encode_test[0], int)
+
+
+@pytest.mark.parametrize("langpair", ["en-de"])
+@pytest.mark.parametrize("source_lines", [["Der Bau und die Reparatur der Autostraßen...", "die Mitteilungen sollen den geschäftlichen kommerziellen Charakter tragen."]])
+@pytest.mark.parametrize("target_lines", [["Construction and repair of highways and...", "An announcement must be commercial character."]]
+def test_collate(langpair, source_lines, target_lines):
+    ds = WMT14dataset(langpair, source_lines, target_lines)
+    source_encode_pad_test, target_encode_pad_test = ds.collate(source_lines[0], target_lines[0])
+    assert source_encode_pad_test.size() == target_encode_pad_test.size()
+    assert source_encode_pad_test.size()[0] == ds.model_config.max_len

tests/test_load_dataset.py

@@ -0,0 +1,9 @@
+import pytest
+
+from src.load_dataset import WMT14DataModule
+
+
+@pytest.mark.parametrize("langpair", ["de-en"])
+def test_setup(langpair):
+    dm = WMT14DataModule(langpair)
+

tests/test_utils.py

@@ -0,0 +1,15 @@
+import pytest  # noqa: F401
+
+from src.utils import read_lines
+
+test_filepath = 'data/example.de'
+
+
+@pytest.mark.parametrize("filepath", test_filepath)
+def test_read_lines(filepath):
+    de = read_lines(filepath)
+    assert isinstance(de, list)
+    assert (
+        de[0]
+        == "iron cement ist eine gebrauchs-fertige Paste, die mit einem Spachtel oder den Fingern als Hohlkehle in die Formecken (Winkel) der Stahlguss -Kokille aufgetragen wird."
+    )