| """Coherence tests for the CORTEX model, tokenizer and checkpoints. |
| |
| These tests exercise real code paths: they build real models, run real forward and |
| backward passes, write real checkpoints to disk and read them back. Nothing is mocked. |
| |
| Run with:: |
| |
| PYTHONPATH=. python3 -m pytest tests/test_cortex_pipeline.py -q |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import pytest |
| import torch |
|
|
| from model.cortex_model import CortexConfig, CortexForCausalLM, count_parameters |
| from model.init import ( |
| CHECKPOINT_FORMAT, |
| init_cortex_model, |
| load_cortex_checkpoint, |
| save_cortex_checkpoint, |
| sha256_file, |
| ) |
| from tokenizer.cortex_tokenizer import ( |
| assert_vocab_matches_config, |
| describe_tokenizer, |
| load_tokenizer, |
| ) |
|
|
| REPO_ROOT = Path(__file__).resolve().parent.parent |
|
|
|
|
| @pytest.fixture(scope="module") |
| def small_config() -> CortexConfig: |
| """A tiny model that trains in milliseconds but exercises every code path.""" |
| return CortexConfig( |
| vocab_size=129280, hidden_size=64, num_hidden_layers=2, |
| num_attention_heads=4, num_key_value_heads=2, intermediate_size=128, |
| max_position_embeddings=64, model_name="cortex-test", |
| ) |
|
|
|
|
| |
|
|
|
|
| def test_parameter_count_matches_analytical(small_config): |
| model = CortexForCausalLM(small_config) |
| assert count_parameters(model) == small_config.parameter_breakdown()["total_estimated"] |
|
|
|
|
| def test_tied_embeddings_share_storage(small_config): |
| model = CortexForCausalLM(small_config) |
| assert small_config.tie_word_embeddings |
| assert model.lm_head.weight is model.embed_tokens.weight |
|
|
|
|
| def test_untied_embeddings_are_separate(small_config): |
| cfg = CortexConfig(**{**small_config.__dict__, "tie_word_embeddings": False}) |
| model = CortexForCausalLM(cfg) |
| assert model.lm_head.weight is not model.embed_tokens.weight |
| |
| assert count_parameters(model) > small_config.parameter_breakdown()["total_estimated"] |
|
|
|
|
| def test_forward_shapes(small_config): |
| model = CortexForCausalLM(small_config) |
| ids = torch.randint(0, small_config.vocab_size, (2, 16)) |
| out = model(ids) |
| assert out["logits"].shape == (2, 16, small_config.vocab_size) |
| assert out["loss"] is None |
|
|
|
|
| def test_loss_is_finite_and_backward_works(small_config): |
| model = CortexForCausalLM(small_config) |
| ids = torch.randint(0, small_config.vocab_size, (2, 16)) |
| out = model(ids, labels=ids) |
| assert torch.isfinite(out["loss"]) |
| out["loss"].backward() |
| grads = [p.grad for p in model.parameters() if p.grad is not None] |
| assert grads and all(torch.isfinite(g).all() for g in grads) |
|
|
|
|
| def test_causal_mask_hides_future_tokens(small_config): |
| """Changing a future token must not change an earlier position's logits.""" |
| model = CortexForCausalLM(small_config).eval() |
| ids = torch.randint(0, small_config.vocab_size, (1, 8)) |
| modified = ids.clone() |
| modified[0, -1] = (modified[0, -1] + 1) % small_config.vocab_size |
| with torch.no_grad(): |
| a = model(ids)["logits"][0, :-1] |
| b = model(modified)["logits"][0, :-1] |
| assert torch.allclose(a, b, atol=1e-5) |
|
|
|
|
| def test_padding_mask_blocks_padded_positions(small_config): |
| model = CortexForCausalLM(small_config).eval() |
| ids = torch.randint(0, small_config.vocab_size, (1, 8)) |
| mask = torch.tensor([[1, 1, 1, 1, 1, 0, 0, 0]]) |
| with torch.no_grad(): |
| out = model(ids, attention_mask=mask)["logits"] |
| assert torch.isfinite(out).all() |
|
|
|
|
| def test_invalid_config_is_rejected(): |
| with pytest.raises(ValueError): |
| CortexConfig(hidden_size=65, num_attention_heads=5, num_key_value_heads=5).validate() |
| with pytest.raises(ValueError): |
| CortexConfig(hidden_size=64, num_attention_heads=6, num_key_value_heads=4).validate() |
|
|
|
|
| def test_greedy_generation_stays_in_vocab(small_config): |
| model = CortexForCausalLM(small_config).eval() |
| ids = torch.randint(0, small_config.vocab_size, (1, 4)) |
| out = model.generate(ids, max_new_tokens=5, temperature=0.0) |
| assert out.shape == (1, 9) |
| assert int(out.max()) < small_config.vocab_size |
|
|
|
|
| def test_config_roundtrip(tmp_path, small_config): |
| path = tmp_path / "config.json" |
| small_config.to_json(path) |
| assert CortexConfig.from_json(path) == small_config |
|
|
|
|
| |
|
|
|
|
| def test_tokenizer_vocab_matches_cortex_config(): |
| tokenizer = load_tokenizer() |
| assert_vocab_matches_config(tokenizer, 129280) |
|
|
|
|
| def test_tokenizer_counters_are_consistent(): |
| facts = describe_tokenizer() |
| assert facts["base_vocab_size"] == 128000 |
| assert facts["added_tokens"] == 1283 |
| |
| assert facts["overlapping_ids"] == 3 |
| assert facts["unique_ids"] == 129280 |
| assert facts["base_vocab_size"] + facts["added_tokens"] - facts["overlapping_ids"] == 129280 |
|
|
|
|
| def test_tokenizer_roundtrip_preserves_code(): |
| tokenizer = load_tokenizer() |
| code = "def add(a, b):\n return a + b\n" |
| assert tokenizer.decode(tokenizer.encode(code).ids) == code |
|
|
|
|
| def test_vocab_mismatch_is_rejected(): |
| tokenizer = load_tokenizer() |
| with pytest.raises(ValueError, match="vocab_size"): |
| assert_vocab_matches_config(tokenizer, 1000) |
|
|
|
|
| |
|
|
|
|
| def test_checkpoint_roundtrip_preserves_weights(tmp_path, small_config): |
| model = init_cortex_model(small_config, seed=7) |
| ckpt = save_cortex_checkpoint(model, small_config, tmp_path / "ckpt", step=3) |
|
|
| loaded, cfg, provenance = load_cortex_checkpoint(ckpt) |
| assert cfg == small_config |
| assert provenance["step"] == 3 |
| assert provenance["format"] == CHECKPOINT_FORMAT |
| for (n1, p1), (n2, p2) in zip( |
| model.state_dict().items(), loaded.state_dict().items() |
| ): |
| assert n1 == n2 |
| assert torch.equal(p1, p2) |
|
|
|
|
| def test_checkpoint_stores_tied_weights_once(tmp_path, small_config): |
| """A tied embedding must be written once, not duplicated, and still load back.""" |
| model = init_cortex_model(small_config, seed=5) |
| ckpt = save_cortex_checkpoint(model, small_config, tmp_path / "ckpt") |
| provenance = json.loads((ckpt / "provenance.json").read_text(encoding="utf-8")) |
| assert provenance["tied_tensors"] == {"lm_head.weight": "embed_tokens.weight"} |
| assert provenance["stored_tensors"] < provenance["tensor_count"] |
|
|
| loaded, _, _ = load_cortex_checkpoint(ckpt) |
| assert torch.equal(loaded.embed_tokens.weight, loaded.lm_head.weight) |
| assert torch.equal(loaded.embed_tokens.weight, model.embed_tokens.weight) |
|
|
|
|
| def test_checkpoint_records_no_false_base_model(tmp_path, small_config): |
| """A from-scratch CORTEX checkpoint must not claim a parent model.""" |
| model = init_cortex_model(small_config, seed=1) |
| ckpt = save_cortex_checkpoint(model, small_config, tmp_path / "ckpt") |
| provenance = json.loads((ckpt / "provenance.json").read_text(encoding="utf-8")) |
| assert provenance["base_model"] is None |
| assert provenance["developer"] == "Frankenstein-Labs" |
| assert "trained by Frankenstein-Labs" in provenance["weights"] |
|
|
|
|
| def test_checkpoint_manifest_lists_every_file(tmp_path, small_config): |
| model = init_cortex_model(small_config, seed=1) |
| ckpt = save_cortex_checkpoint(model, small_config, tmp_path / "ckpt") |
| manifest = json.loads((ckpt / "manifest.json").read_text(encoding="utf-8")) |
| on_disk = {p.name for p in ckpt.iterdir() if p.name != "manifest.json"} |
| assert set(manifest["files"]) == on_disk |
| for name, info in manifest["files"].items(): |
| assert info["sha256"] == sha256_file(ckpt / name) |
|
|
|
|
| def test_corrupted_checkpoint_is_detected(tmp_path, small_config): |
| model = init_cortex_model(small_config, seed=1) |
| ckpt = save_cortex_checkpoint(model, small_config, tmp_path / "ckpt") |
| weights = ckpt / "weights.safetensors" |
| if not weights.exists(): |
| pytest.skip("safetensors not installed; pointer fallback in use") |
| raw = bytearray(weights.read_bytes()) |
| raw[-1] ^= 0xFF |
| weights.write_bytes(bytes(raw)) |
| with pytest.raises(ValueError, match="sha256 mismatch"): |
| load_cortex_checkpoint(ckpt) |
|
|
|
|
| def test_missing_manifest_is_rejected(tmp_path, small_config): |
| empty = tmp_path / "not-a-checkpoint" |
| empty.mkdir() |
| with pytest.raises(FileNotFoundError, match="manifest.json"): |
| load_cortex_checkpoint(empty) |
|
|
|
|
| def test_checkpoint_is_reproducible_from_seed(tmp_path, small_config): |
| a = init_cortex_model(small_config, seed=42) |
| b = init_cortex_model(small_config, seed=42) |
| for (_, p1), (_, p2) in zip(a.state_dict().items(), b.state_dict().items()): |
| assert torch.equal(p1, p2) |
|
|
|
|
| def test_different_seeds_give_different_weights(small_config): |
| a = init_cortex_model(small_config, seed=1) |
| b = init_cortex_model(small_config, seed=2) |
| assert not torch.equal(a.embed_tokens.weight, b.embed_tokens.weight) |
|
|
|
|
| |
|
|
|
|
| def test_training_refuses_undeclared_dataset(): |
| from training.train import assert_dataset_usable, load_registry |
|
|
| registry = load_registry() |
| with pytest.raises(KeyError, match="not declared"): |
| assert_dataset_usable(registry, "does-not-exist") |
|
|
|
|
| def test_training_refuses_unlicensed_dataset(): |
| from training.train import assert_dataset_usable |
|
|
| registry = {"datasets": [{ |
| "id": "bad", "source": "x", "license": "unknown", |
| "license_url": "x", "origin": "x", "redistribution": "no", |
| }]} |
| with pytest.raises(ValueError, match="refusing to train"): |
| assert_dataset_usable(registry, "bad") |
|
|
|
|
| def test_every_registered_dataset_has_license_and_origin(): |
| from training.train import load_registry |
|
|
| for entry in load_registry()["datasets"]: |
| assert entry["license"], entry["id"] |
| assert entry["origin"], entry["id"] |
| assert entry["license_url"], entry["id"] |
| assert entry["license"].lower() not in {"unknown", "none", "proprietary"}, entry["id"] |
|
|
|
|
| def test_dataset_metadata_files_match_their_jsonl(): |
| from training.train import load_registry |
|
|
| for entry in load_registry()["datasets"]: |
| jsonl = REPO_ROOT / entry["path"] |
| meta_path = jsonl.with_suffix(".meta.json") |
| if not jsonl.exists() or not meta_path.exists(): |
| continue |
| meta = json.loads(meta_path.read_text(encoding="utf-8")) |
| assert meta["jsonl_sha256"] == sha256_file(jsonl), entry["id"] |
| assert meta["license"] == entry["license"], entry["id"] |
|
|
|
|
| def test_build_token_stream_produces_full_windows(small_config): |
| from training.train import build_token_stream |
|
|
| tokenizer = load_tokenizer() |
| records = [{"text": "def f():\n return 1\n"}] * 40 |
| data = build_token_stream(records, tokenizer, seq_len=32) |
| assert data.ndim == 2 |
| assert data.shape[1] == 33 |
| assert int(data.max()) < small_config.vocab_size |
|
|
|
|
| def test_training_step_reduces_loss_on_a_tiny_corpus(): |
| """Overfit a tiny corpus: if the pipeline works, loss must drop sharply.""" |
| from training.train import CortexTrainer, build_token_stream |
|
|
| cfg = CortexConfig( |
| vocab_size=129280, hidden_size=64, num_hidden_layers=2, |
| num_attention_heads=4, num_key_value_heads=2, intermediate_size=128, |
| max_position_embeddings=64, |
| ) |
| tokenizer = load_tokenizer() |
| records = [{"text": "def add(a, b):\n return a + b\n"}] * 64 |
| data = build_token_stream(records, tokenizer, seq_len=32) |
| model = init_cortex_model(cfg, seed=0) |
| train_cfg = { |
| "micro_batch_size": 2, "gradient_accumulation_steps": 1, |
| "learning_rate": 3e-3, "min_learning_rate": 3e-4, "weight_decay": 0.0, |
| "beta1": 0.9, "beta2": 0.95, "grad_clip": 1.0, |
| "warmup_steps": 1, "max_steps": 40, "seed": 0, |
| "log_every": 1, "save_every": 0, "sequence_length": 32, |
| } |
| trainer = CortexTrainer(model, cfg, tokenizer, data, train_cfg, device="cpu") |
| history = trainer.train(max_steps=40, log_every=1) |
| first = sum(h["loss"] for h in history[:3]) / 3 |
| last = sum(h["loss"] for h in history[-3:]) / 3 |
| assert last < first, f"loss did not fall: {first} -> {last}" |
|
|
