tests/test_model.py

"""
Test suite for the Dragoman word alignment model.

Run with: python -m pytest tests/ -x -v
Skip slow tests: python -m pytest tests/ -x -v -m "not slow"
"""

import json
import struct
from pathlib import Path

import pytest

MODEL_DIR = Path(__file__).resolve().parent.parent

# ---------------------------------------------------------------------------
# 1. Model card validation
# ---------------------------------------------------------------------------

class TestModelCard:
    """Validate README.md model card and YAML frontmatter."""

    def test_readme_exists(self):
        assert (MODEL_DIR / "README.md").is_file()

    def test_yaml_frontmatter_present(self):
        text = (MODEL_DIR / "README.md").read_text()
        assert text.startswith("---"), "README should start with YAML frontmatter"
        # Must have closing ---
        second = text.index("---", 3)
        assert second > 3

    def test_yaml_required_fields(self):
        text = (MODEL_DIR / "README.md").read_text()
        end = text.index("---", 3)
        frontmatter = text[3:end]
        for field in ("language:", "license:", "base_model:", "pipeline_tag:"):
            assert field in frontmatter, f"Missing frontmatter field: {field}"

    def test_yaml_languages(self):
        text = (MODEL_DIR / "README.md").read_text()
        end = text.index("---", 3)
        frontmatter = text[3:end]
        for lang in ("grc", "el", "en"):
            assert lang in frontmatter, f"Missing language: {lang}"


# ---------------------------------------------------------------------------
# 2. Config integrity
# ---------------------------------------------------------------------------

class TestConfig:
    """Validate config.json structure."""

    @pytest.fixture(scope="class")
    def config(self):
        with open(MODEL_DIR / "config.json") as f:
            return json.load(f)

    def test_config_valid_json(self, config):
        assert isinstance(config, dict)

    def test_architecture(self, config):
        assert "architectures" in config
        assert "XLMRobertaModel" in config["architectures"]

    def test_model_type(self, config):
        assert config.get("model_type") == "xlm-roberta"

    def test_hidden_size(self, config):
        assert config.get("hidden_size") == 768

    def test_num_layers(self, config):
        assert config.get("num_hidden_layers") == 12

    def test_vocab_size(self, config):
        assert config.get("vocab_size") == 250002


# ---------------------------------------------------------------------------
# 3. Tokenizer integrity
# ---------------------------------------------------------------------------

class TestTokenizer:
    """Validate tokenizer files and basic encoding."""

    def test_tokenizer_json_exists(self):
        assert (MODEL_DIR / "tokenizer.json").is_file()

    def test_tokenizer_config_exists(self):
        assert (MODEL_DIR / "tokenizer_config.json").is_file()

    def test_tokenizer_json_valid(self):
        with open(MODEL_DIR / "tokenizer.json") as f:
            data = json.load(f)
        assert "model" in data

    def test_tokenizer_config_valid(self):
        with open(MODEL_DIR / "tokenizer_config.json") as f:
            data = json.load(f)
        assert data.get("tokenizer_class") == "XLMRobertaTokenizer"

    def test_encode_polytonic_greek(self):
        """Tokenizer should handle polytonic Ancient Greek without errors."""
        from tokenizers import Tokenizer

        tok = Tokenizer.from_file(str(MODEL_DIR / "tokenizer.json"))
        enc = tok.encode("μῆνιν ἄειδε θεὰ Πηληϊάδεω Ἀχιλῆος")
        assert len(enc.ids) > 0
        assert len(enc.tokens) > 0

    def test_encode_modern_greek(self):
        """Tokenizer should handle Modern Greek."""
        from tokenizers import Tokenizer

        tok = Tokenizer.from_file(str(MODEL_DIR / "tokenizer.json"))
        enc = tok.encode("Ψάλλε θεά τον τρομερό θυμό του Αχιλλέα")
        assert len(enc.ids) > 0

    def test_encode_english(self):
        """Tokenizer should handle English."""
        from tokenizers import Tokenizer

        tok = Tokenizer.from_file(str(MODEL_DIR / "tokenizer.json"))
        enc = tok.encode("Sing O goddess the wrath of Achilles")
        assert len(enc.ids) > 0

    def test_greek_tokens_nonempty(self):
        """Each Greek word should produce at least one non-special token."""
        from tokenizers import Tokenizer

        tok = Tokenizer.from_file(str(MODEL_DIR / "tokenizer.json"))
        words = ["μῆνιν", "ἄειδε", "θεὰ", "Πηληϊάδεω", "Ἀχιλῆος"]
        for word in words:
            enc = tok.encode(word)
            # At least 1 token per word
            assert len(enc.ids) >= 1, f"No tokens for '{word}'"


# ---------------------------------------------------------------------------
# 4. Model weights
# ---------------------------------------------------------------------------

class TestModelWeights:
    """Validate model.safetensors file."""

    def test_safetensors_exists(self):
        assert (MODEL_DIR / "model.safetensors").is_file()

    def test_safetensors_size(self):
        size_mb = (MODEL_DIR / "model.safetensors").stat().st_size / (1024 * 1024)
        # XLM-R base is ~1.1 GB in fp32
        assert size_mb > 500, f"model.safetensors too small: {size_mb:.0f} MB"

    def test_safetensors_header(self):
        """Safetensors file should have a valid header."""
        with open(MODEL_DIR / "model.safetensors", "rb") as f:
            header_len = struct.unpack("<Q", f.read(8))[0]
            assert 0 < header_len < 10_000_000, f"Bad header length: {header_len}"
            header_bytes = f.read(header_len)
            header = json.loads(header_bytes)
            assert isinstance(header, dict)
            # Should contain embedding keys
            keys = [k for k in header if k != "__metadata__"]
            assert len(keys) > 50, f"Only {len(keys)} tensors in safetensors"

    @pytest.mark.slow
    def test_load_safetensors(self):
        """Load all tensors from safetensors file."""
        from safetensors import safe_open

        with safe_open(MODEL_DIR / "model.safetensors", framework="pt") as f:
            keys = f.keys()
            assert len(keys) > 50
            # Spot-check one tensor
            emb = f.get_tensor("embeddings.word_embeddings.weight")
            assert emb.shape[0] == 250002
            assert emb.shape[1] == 768


# ---------------------------------------------------------------------------
# 5. Lemma head
# ---------------------------------------------------------------------------

class TestLemmaHead:
    """Validate lemma_head.pt checkpoint."""

    def test_lemma_head_exists(self):
        assert (MODEL_DIR / "lemma_head.pt").is_file()

    def test_lemma_head_loadable(self):
        import torch

        data = torch.load(MODEL_DIR / "lemma_head.pt", map_location="cpu", weights_only=True)
        assert isinstance(data, dict), "lemma_head.pt should contain a state dict"
        assert len(data) > 0, "state dict is empty"

    def test_lemma_head_shapes(self):
        """Lemma head should have layers compatible with hidden_size=768."""
        import torch

        data = torch.load(MODEL_DIR / "lemma_head.pt", map_location="cpu", weights_only=True)
        # Check that at least one weight has 768 as a dimension
        has_768 = any(
            768 in t.shape for t in data.values() if hasattr(t, "shape")
        )
        assert has_768, "No tensor with hidden_size=768 found in lemma head"


# ---------------------------------------------------------------------------
# 6. Basic inference (slow)
# ---------------------------------------------------------------------------

class TestInference:
    """End-to-end inference tests. Require transformers + torch."""

    @pytest.mark.slow
    def test_model_loads(self):
        """Model loads via transformers AutoModel."""
        from transformers import AutoModel

        model = AutoModel.from_pretrained(str(MODEL_DIR))
        assert model.config.hidden_size == 768

    @pytest.mark.slow
    def test_simalign_alignment(self):
        """Run a simple AG-EN alignment via simalign."""
        from simalign import SentenceAligner

        aligner = SentenceAligner(
            model=str(MODEL_DIR),
            token_type="bpe",
            matching_methods="i",
            device="cpu",
            layer=8,
        )
        ag = ["μῆνιν", "ἄειδε", "θεὰ", "Πηληϊάδεω", "Ἀχιλῆος"]
        en = ["Sing", "O", "goddess", "the", "wrath", "of", "Achilles"]
        result = aligner.get_word_aligns(ag, en)

        pairs = result["itermax"]
        assert isinstance(pairs, list)
        assert len(pairs) > 0

        # Each pair should be (int, int) within bounds
        for a, b in pairs:
            assert 0 <= a < len(ag)
            assert 0 <= b < len(en)

        # Sanity: Ἀχιλῆος (idx 4) should align to Achilles (idx 6)
        aligned_targets = {a: b for a, b in pairs}
        assert aligned_targets.get(4) == 6, (
            f"Expected Achilles alignment (4->6), got {aligned_targets}"
        )