File size: 3,254 Bytes
eca5751 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | """Tests for Nexus Coder model."""
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import pytest
import torch
from nexus.config import NexusConfig
from nexus.model.nexus_coder import NexusCoderForCausalLM
from nexus.tokenizer.tokenizer import NexusTokenizer
from nexus.training.dataset import NexusDataset, AUTHOR_TRAINING_DATA
@pytest.fixture
def tiny_config():
return NexusConfig(
vocab_size=500,
hidden_size=64,
num_hidden_layers=2,
num_attention_heads=4,
num_kv_heads=2,
head_dim=16,
intermediate_size=128,
num_experts=4,
num_active_experts=2,
max_position_embeddings=128,
)
@pytest.fixture
def tiny_model(tiny_config):
return NexusCoderForCausalLM(tiny_config)
def test_config_default():
"""Test default config."""
config = NexusConfig()
assert config.hidden_size == 2048
assert config.num_hidden_layers == 12
assert config.num_experts == 24
assert config.num_active_experts == 3
assert config.max_position_embeddings == 50000
def test_param_count():
"""Test parameter count is ~10B / 1.5B."""
config = NexusConfig()
stats = config.estimated_total_params()
assert 9.5e9 < stats["total_params"] < 11e9
assert 1.3e9 < stats["active_params"] < 1.7e9
def test_model_forward(tiny_model):
"""Test model forward pass."""
input_ids = torch.randint(0, 500, (2, 16))
outputs = tiny_model(input_ids=input_ids)
assert outputs["logits"].shape == (2, 16, 500)
def test_model_training(tiny_model):
"""Test model with labels (training)."""
input_ids = torch.randint(0, 500, (2, 16))
labels = input_ids.clone()
outputs = tiny_model(input_ids=input_ids, labels=labels)
assert outputs["loss"] is not None
assert outputs["loss"].item() > 0
def test_generate(tiny_model):
"""Test generation."""
input_ids = torch.randint(0, 500, (1, 4))
generated = tiny_model.generate(
input_ids=input_ids,
max_new_tokens=5,
do_sample=False,
)
assert generated.shape[0] == 1
assert generated.shape[1] >= 4
def test_tokenizer():
"""Test tokenizer basic."""
tokenizer = NexusTokenizer(vocab_size=1000)
corpus = ["hello world nexus coder hieu louis"]
tokenizer.train(corpus)
ids = tokenizer.encode("hello nexus")
assert len(ids) > 0
decoded = tokenizer.decode(ids)
assert "hello" in decoded.lower() or "nexus" in decoded.lower()
def test_dataset():
"""Test dataset."""
assert len(AUTHOR_TRAINING_DATA) > 0
# Check author info is present
info_texts = " ".join([
f"{d['system']} {d['user']} {d['assistant']}" for d in AUTHOR_TRAINING_DATA
])
assert "Hieu Louis" in info_texts
assert "2026" in info_texts
def test_author_info_hardcoded():
"""Test that author info is hardcoded in dataset."""
from nexus.training.dataset import get_author_info
info = get_author_info()
assert info["name"] == "Hieu Louis"
assert info["github"] == "mhieuhonda"
assert info["year"] == "2026"
assert info["model_name"] == "Nexus Coder"
if __name__ == "__main__":
pytest.main([__file__, "-v"])
|