tests/test_cortex_pipeline.py

"""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",
    )


# --------------------------------------------------------------------------- model


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
    # two embedding matrices now, so the count must be larger
    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


# ----------------------------------------------------------------------- tokenizer


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
    # ids 0, 1, 2 appear in both the base vocab and the added tokens
    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)


# --------------------------------------------------------------------- checkpoints


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)


# ------------------------------------------------------------------- training data


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