File size: 4,575 Bytes
b88c26d | 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 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | """
BURT-IMMA pytest configuration
License: BSL-1.1
Contact: jessica@collectivekitty.com
"""
import os
import pytest
import numpy as np
try:
import torch
HAS_TORCH = True
except ImportError:
HAS_TORCH = False
try:
import yaml
HAS_YAML = True
except ImportError:
HAS_YAML = False
# ---------------------------------------------------------------------------
# Custom marks
# ---------------------------------------------------------------------------
def pytest_configure(config):
config.addinivalue_line("markers", "slow: marks tests as slow (deselect with '-m \"not slow\"')")
config.addinivalue_line("markers", "cuda: marks tests requiring CUDA GPU")
config.addinivalue_line("markers", "integration: marks integration tests")
def pytest_collection_modifyitems(config, items):
"""Skip CUDA tests if no GPU available."""
skip_cuda = pytest.mark.skip(reason="CUDA not available")
for item in items:
if "cuda" in item.keywords:
if not HAS_TORCH or not torch.cuda.is_available():
item.add_marker(skip_cuda)
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture
def device():
"""Return device string: 'cuda' if available, else 'cpu'."""
if HAS_TORCH and torch.cuda.is_available():
return "cuda"
return "cpu"
@pytest.fixture
def config():
"""Load ablation_arithmetic.yaml config."""
config_path = os.path.join(
os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
"config",
"ablation_arithmetic.yaml"
)
if HAS_YAML and os.path.exists(config_path):
with open(config_path, "r") as f:
return yaml.safe_load(f)
else:
# Fallback default config
return {
"experiment": {"name": "mmep_arithmetic_ablation", "seed": 42, "device": "cpu"},
"model": {
"hidden_dim": 256,
"num_layers": 4,
"num_experts": 4,
"top_k": 1,
"d_mem": 256,
},
"mmep": {
"T_free": 20,
"T_nudge": 4,
"alpha_relax": 0.5,
"beta": 0.1,
"rho_ret": 1.0,
"rho_inst": 0.5,
"lambda_max": 0.95,
"lr_W": 0.001,
"lr_C_global": 0.01,
"lr_C_expert": 0.005,
},
"data": {
"corpus_size": 10000,
"train_queries": 5000,
"val_queries": 500,
"test_queries": 1000,
"max_seq_len": 128,
},
"training": {
"batch_size": 32,
"num_epochs": 50,
"eval_every": 5,
"checkpoint_every": 10,
"gradient_clip": 1.0,
},
"ablations": [
{"name": "full_mmep", "description": "Complete MMEP", "disable": []},
{"name": "no_memory", "disable": ["C_global", "C_expert"]},
{"name": "no_constraint", "disable": ["projection"]},
{"name": "no_moe", "override": {"num_experts": 1}},
{"name": "standard_bp", "disable": ["mmep"], "use_backprop": True},
],
"logging": {
"wandb_project": "burt-imma-ablation",
"log_entropy": True,
"log_spectral_norm": True,
"log_memory_trace": True,
},
}
@pytest.fixture
def random_seed():
"""Set deterministic random seed (42) for reproducibility."""
seed = 42
np.random.seed(seed)
if HAS_TORCH:
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
return seed
@pytest.fixture
def sample_batch(device, config, random_seed):
"""Generate a random tensor batch for testing."""
batch_size = config["training"]["batch_size"]
hidden_dim = config["model"]["hidden_dim"]
seq_len = 16 # Short sequence for fast tests
if HAS_TORCH:
x = torch.randn(batch_size, seq_len, hidden_dim)
if device == "cuda":
x = x.cuda()
return x
else:
return np.random.randn(batch_size, seq_len, hidden_dim).astype(np.float32)
|