tests/test_tokenizer_fast.py

"""Tests for LatinBertTokenizerFast (Rust-backed, word_ids support).

Guards the two fixes to the fast tokenizer:
  1. whitespace-split before char-class encoding (no inter-word space
     escapes; correct word_ids alignment)
  2. lowercasing normalizer (do_lower_case), so capitalized input is not
     escaped to codepoints — matching the slow tokenizer.
"""

import importlib.util
from pathlib import Path

import pytest

REPO_ROOT = Path(__file__).resolve().parent.parent
VOCAB_FILE = str(Path(__file__).parent / "latin.subword.encoder")


def _load_fast_module():
    path = REPO_ROOT / "tokenization_latin_bert_fast.py"
    spec = importlib.util.spec_from_file_location("tlbf_under_test", path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


@pytest.fixture(scope="module")
def fast_module():
    return _load_fast_module()


@pytest.fixture
def fast(fast_module):
    return fast_module.LatinBertTokenizerFast(vocab_file=VOCAB_FILE)


@pytest.fixture
def slow():
    from latincy_latinbert import LatinBertTokenizer
    return LatinBertTokenizer(vocab_file=VOCAB_FILE)


class TestFastTokenizer:
    def test_no_interword_space_escapes(self, fast):
        enc = fast("gallia est omnis divisa in partes tres")
        toks = fast.convert_ids_to_tokens(enc["input_ids"])
        assert all(t not in ("\\", "32", ";_") for t in toks)

    def test_word_ids_alignment(self, fast):
        enc = fast("Gallia est omnis")
        # [CLS]=None, gallia=0, est=1, omnis=2, [SEP]=None
        assert enc.word_ids() == [None, 0, 1, 2, None]

    def test_lowercases_like_slow(self, fast):
        """Capitalized input must not produce codepoint escapes."""
        enc = fast("Gallia")
        toks = fast.convert_ids_to_tokens(enc["input_ids"])
        assert "gallia_" in toks
        assert all(not t.isdigit() for t in toks)

    def test_matches_slow_ids(self, fast, slow):
        for text in ["Gallia est omnis divisa in partes tres",
                     "arma virumque cano",
                     "ROMA",
                     "gallia est omnis divisa in partes tres."]:
            s = slow.encode(text, add_special_tokens=False)
            f = fast.encode(text, add_special_tokens=False)
            assert s == f, f"slow/fast mismatch for {text!r}: {s} != {f}"