""" 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(" 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}" )