| """ |
| tests/test_collators.py |
| Unit tests for data collators. |
| """ |
|
|
| import sys |
| sys.path.insert(0, "src") |
|
|
| import pytest |
|
|
|
|
| def test_find_last_subseq(): |
| from ankahi.data.collators import _find_last_subseq |
| seq = [1, 2, 3, 1, 2, 3, 4] |
| assert _find_last_subseq(seq, [1, 2]) == 3 |
| assert _find_last_subseq(seq, [3, 4]) == 5 |
| assert _find_last_subseq(seq, [9]) is None |
| assert _find_last_subseq([], [1]) is None |
|
|
|
|
| def test_find_last_subseq_single(): |
| from ankahi.data.collators import _find_last_subseq |
| assert _find_last_subseq([5, 5, 5], [5]) == 2 |
| assert _find_last_subseq([1, 2, 3], [1, 2, 3]) == 0 |
|
|
|
|
| def test_pack_texts_empty(): |
| from ankahi.data.collators import pack_texts |
|
|
| class MockTokenizer: |
| def encode(self, text, add_special_tokens=True): |
| return list(range(len(text.split()))) |
| def decode(self, tokens): |
| return " ".join(str(t) for t in tokens) |
|
|
| packed = pack_texts([], MockTokenizer(), max_length=512) |
| assert packed == [] |
|
|
|
|
| def test_pack_texts_basic(): |
| from ankahi.data.collators import pack_texts |
|
|
| class MockTokenizer: |
| def encode(self, text, add_special_tokens=True): |
| return [0] * 100 |
| def decode(self, tokens): |
| return " ".join(str(t) for t in tokens) |
|
|
| texts = ["text " * 20] * 5 |
| |
| packed = pack_texts(texts, MockTokenizer(), max_length=512) |
| assert len(packed) >= 1 |
|
|
