ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: eliminate test duplication via shared helpers

Browse Source 
- Add tests/helpers.py with shared config, dataset, tokenizer, executor, and assertion helpers
- Replace 15 copies of device one-liner with session-scoped fixture
- Collapse 5 near-identical Dataset subclasses into RandomTokenDataset
- Remove duplicate _make_config/_make_model/_make_frozen and FakeTokenizer/FakeExecutor definitions
- Make test_callbacks and test_early_stopping use existing train_config_factory
- Replace 6 duplicate meta.json read blocks with load_shard_meta
- Fix mkdtemp leaks in test_lora.py with TemporaryDirectory
This commit is contained in:
ViperEkura
2026-07-27 22:34:53 +08:00
parent c26a47b0df
commit 5ba21f4eb3
16 changed files with 364 additions and 661 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/module/test_encoder.py
+23 -43
View File
@@ -9,34 +9,22 @@ import torch
from astrai.config.model_config import EncoderConfig
from astrai.model.automodel import AutoModel
from astrai.model.encoder import EmbeddingEncoder
TINY_CONFIG = dict(
vocab_size=128,
hidden_size=8,
num_attention_heads=2,
num_key_value_heads=1,
intermediate_size=16,
max_position_embeddings=64,
num_hidden_layers=2,
rms_norm_eps=1e-5,
)
_device = "cuda" if torch.cuda.is_available() else "cpu"
from tests.helpers import TINY_CONFIG, assert_state_dicts_equal
def _make_model(**kwargs):
def _make_model(device, **kwargs):
config = EncoderConfig(**{**TINY_CONFIG, **kwargs})
return EmbeddingEncoder(config).to(device=_device)
return EmbeddingEncoder(config).to(device=device)
@pytest.mark.parametrize("pooling_type", ["mean", "cls", "last"])
def test_encoder_forward_pooling(pooling_type):
model = _make_model(pooling_type=pooling_type)
def test_encoder_forward_pooling(pooling_type, device):
model = _make_model(device, pooling_type=pooling_type)
model.eval()
batch_size, seq_len = 2, 8
input_ids = torch.randint(
0, TINY_CONFIG["vocab_size"], (batch_size, seq_len), device=_device
0, TINY_CONFIG["vocab_size"], (batch_size, seq_len), device=device
)
with torch.no_grad():
@@ -46,15 +34,15 @@ def test_encoder_forward_pooling(pooling_type):
assert not torch.isnan(output).any()
def test_encoder_forward_with_padding():
model = _make_model()
def test_encoder_forward_with_padding(device):
model = _make_model(device)
model.eval()
batch_size, seq_len = 2, 8
input_ids = torch.randint(
0, TINY_CONFIG["vocab_size"], (batch_size, seq_len), device=_device
0, TINY_CONFIG["vocab_size"], (batch_size, seq_len), device=device
)
input_mask = torch.ones(batch_size, seq_len, dtype=torch.bool, device=_device)
input_mask = torch.ones(batch_size, seq_len, dtype=torch.bool, device=device)
input_mask[:, 4:] = False
with torch.no_grad():
@@ -64,13 +52,13 @@ def test_encoder_forward_with_padding():
assert not torch.isnan(output).any()
def test_encoder_normalize():
model = _make_model(pooling_type="mean", normalize_embeddings=True)
def test_encoder_normalize(device):
model = _make_model(device, pooling_type="mean", normalize_embeddings=True)
model.eval()
batch_size, seq_len = 2, 8
input_ids = torch.randint(
0, TINY_CONFIG["vocab_size"], (batch_size, seq_len), device=_device
0, TINY_CONFIG["vocab_size"], (batch_size, seq_len), device=device
)
with torch.no_grad():
@@ -86,26 +74,24 @@ def test_encoder_register():
assert cls is EmbeddingEncoder
def test_encoder_from_transformer_checkpoint():
model = _make_model()
def test_encoder_from_transformer_checkpoint(device):
model = _make_model(device)
state_dict = model.state_dict()
state_dict["lm_head.weight"] = torch.randn(
TINY_CONFIG["vocab_size"], TINY_CONFIG["hidden_size"], device=_device
TINY_CONFIG["vocab_size"], TINY_CONFIG["hidden_size"], device=device
)
new_model = _make_model()
new_model = _make_model(device)
new_model.load_state_dict(state_dict, strict=True)
for key in model.state_dict():
assert torch.equal(new_model.state_dict()[key], model.state_dict()[key])
assert_state_dicts_equal(new_model.state_dict(), model.state_dict())
def test_encoder_save_load():
test_dir = tempfile.mkdtemp(prefix="encoder_test_")
config_path = os.path.join(test_dir, "config.json")
weights_path = os.path.join(test_dir, "model.safetensors")
def test_encoder_save_load(device):
with tempfile.TemporaryDirectory(prefix="encoder_test_") as test_dir:
config_path = os.path.join(test_dir, "config.json")
weights_path = os.path.join(test_dir, "model.safetensors")
try:
config_data = {**TINY_CONFIG, "pooling_type": "mean"}
with open(config_path, "w") as f:
json.dump(config_data, f)
@@ -117,10 +103,4 @@ def test_encoder_save_load():
loaded = EmbeddingEncoder(config)
loaded.load_state_dict(st.load_file(weights_path))
for key in original.state_dict():
assert torch.equal(original.state_dict()[key], loaded.state_dict()[key])
finally:
if os.path.exists(test_dir):
for f in os.listdir(test_dir):
os.remove(os.path.join(test_dir, f))
os.rmdir(test_dir)
assert_state_dicts_equal(original.state_dict(), loaded.state_dict())
tests/module/test_forward_configs.py
+7 -17
View File
@@ -1,20 +1,8 @@
import pytest
import torch
from astrai.config.model_config import AutoRegressiveLMConfig
from astrai.model.transformer import AutoRegressiveLM
TINY_CONFIG = dict(
vocab_size=128,
hidden_size=8,
num_attention_heads=2,
num_key_value_heads=1,
intermediate_size=16,
max_position_embeddings=64,
num_hidden_layers=2,
rms_norm_eps=1e-5,
)
from tests.helpers import TINY_CONFIG
CONFIGS = [
pytest.param(
@@ -70,9 +58,10 @@ CONFIGS = [
@pytest.mark.parametrize("config_kwargs", CONFIGS)
def test_model_forward(config_kwargs):
def test_model_forward(config_kwargs, device):
from astrai.config.model_config import AutoRegressiveLMConfig
config = AutoRegressiveLMConfig(**config_kwargs)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = AutoRegressiveLM(config).to(device=device)
model.eval()
@@ -97,9 +86,10 @@ def test_model_forward(config_kwargs):
@pytest.mark.parametrize("config_kwargs", CONFIGS)
def test_model_forward_with_padding(config_kwargs):
def test_model_forward_with_padding(config_kwargs, device):
from astrai.config.model_config import AutoRegressiveLMConfig
config = AutoRegressiveLMConfig(**config_kwargs)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = AutoRegressiveLM(config).to(device=device)
model.eval()
tests/module/test_lora.py
+28 -29
View File
@@ -249,17 +249,17 @@ def test_save_load_roundtrip():
with torch.no_grad():
out_src = model(x)["logits"].clone()
tmpdir = tempfile.mkdtemp()
save_lora(model, tmpdir, cfg)
with tempfile.TemporaryDirectory() as tmpdir:
save_lora(model, tmpdir, cfg)
model2 = _make_model()
model2.load_state_dict(model.state_dict(), strict=False)
load_lora(model2, tmpdir)
model2 = _make_model()
model2.load_state_dict(model.state_dict(), strict=False)
load_lora(model2, tmpdir)
with torch.no_grad():
out_dst = model2(x)["logits"]
with torch.no_grad():
out_dst = model2(x)["logits"]
torch.testing.assert_close(out_src, out_dst)
torch.testing.assert_close(out_src, out_dst)
def test_save_after_merge_raises():
@@ -271,13 +271,13 @@ def test_save_after_merge_raises():
if isinstance(m, LoRALinear):
m.lora_B.fill_(0.5)
tmpdir = tempfile.mkdtemp()
save_lora(model, tmpdir, cfg)
merge_lora(model)
with tempfile.TemporaryDirectory() as tmpdir:
save_lora(model, tmpdir, cfg)
merge_lora(model)
tmpdir2 = tempfile.mkdtemp()
with pytest.raises(RuntimeError, match="No LoRA parameters"):
save_lora(model, tmpdir2, cfg)
with tempfile.TemporaryDirectory() as tmpdir2:
with pytest.raises(RuntimeError, match="No LoRA parameters"):
save_lora(model, tmpdir2, cfg)
def test_load_lora_on_already_injected():
@@ -289,16 +289,15 @@ def test_load_lora_on_already_injected():
if isinstance(m, LoRALinear):
m.lora_B.fill_(0.5)
tmpdir = tempfile.mkdtemp()
save_lora(model, tmpdir, LoRAConfig(r=4, alpha=8, target_modules=("q_proj",)))
with tempfile.TemporaryDirectory() as tmpdir:
save_lora(model, tmpdir, LoRAConfig(r=4, alpha=8, target_modules=("q_proj",)))
model2 = _make_model()
model2.load_state_dict(model.state_dict(), strict=False)
inject_lora(model2, r=4, alpha=8, target_modules={"q_proj"})
model2 = _make_model()
model2.load_state_dict(model.state_dict(), strict=False)
inject_lora(model2, r=4, alpha=8, target_modules={"q_proj"})
# load onto already-injected model
load_lora(model2, tmpdir)
assert _get_lora_count(model2) > 0
load_lora(model2, tmpdir)
assert _get_lora_count(model2) > 0
def test_load_lora_mismatched_r_raises():
@@ -310,15 +309,15 @@ def test_load_lora_mismatched_r_raises():
if isinstance(m, LoRALinear):
m.lora_B.fill_(0.5)
tmpdir = tempfile.mkdtemp()
save_lora(model, tmpdir, cfg)
with tempfile.TemporaryDirectory() as tmpdir:
save_lora(model, tmpdir, cfg)
model2 = _make_model()
model2.load_state_dict(model.state_dict(), strict=False)
inject_lora(model2, r=4, alpha=8, target_modules={"q_proj"})
model2 = _make_model()
model2.load_state_dict(model.state_dict(), strict=False)
inject_lora(model2, r=4, alpha=8, target_modules={"q_proj"})
with pytest.raises(RuntimeError, match="size mismatch"):
load_lora(model2, tmpdir) # strict=False, only lora keys
with pytest.raises(RuntimeError, match="size mismatch"):
load_lora(model2, tmpdir)
def test_merge_preserves_output():
tests/module/test_tie_weight.py
+7 -42
View File
@@ -1,6 +1,5 @@
import json
import os
import tempfile
import pytest
import safetensors.torch as st
@@ -8,43 +7,13 @@ import torch
from astrai.config.model_config import AutoRegressiveLMConfig
from astrai.model.transformer import AutoRegressiveLM
from tests.helpers import TINY_CONFIG
@pytest.fixture
def transformer_test_env():
test_dir = tempfile.mkdtemp(prefix="transformer_test_")
config_path = os.path.join(test_dir, "config.json")
def test_tie_weight_init(base_test_env):
config_path = base_test_env["config_path"]
config = {
"vocab_size": 1000,
"hidden_size": 8,
"num_attention_heads": 2,
"num_key_value_heads": 1,
"intermediate_size": 16,
"max_position_embeddings": 64,
"num_hidden_layers": 2,
"rms_norm_eps": 1e-5,
}
with open(config_path, "w") as f:
json.dump(config, f)
yield {"test_dir": test_dir, "config_path": config_path, "config": config}
if os.path.exists(test_dir):
try:
for file in os.listdir(test_dir):
os.remove(os.path.join(test_dir, file))
os.rmdir(test_dir)
except Exception:
pass
def test_tie_weight_init(transformer_test_env):
config_path = transformer_test_env["config_path"]
config_data = transformer_test_env["config"].copy()
# case 1: tie weight
config_data = TINY_CONFIG.copy()
config_data["tie_word_embeddings"] = True
with open(config_path, "w") as f:
@@ -62,7 +31,6 @@ def test_tie_weight_init(transformer_test_env):
assert torch.equal(model.lm_head.weight, model.embed_tokens.weight)
assert not torch.equal(model.lm_head.weight, original_weight)
# case 2: not tie weight
config_data["tie_word_embeddings"] = False
with open(config_path, "w") as f:
@@ -81,13 +49,11 @@ def test_tie_weight_init(transformer_test_env):
assert not torch.equal(model.lm_head.weight, original_weight)
def test_model_save_load_with_tie_weight(transformer_test_env):
test_dir = transformer_test_env["test_dir"]
def test_model_save_load_with_tie_weight(base_test_env):
test_dir = base_test_env["test_dir"]
model_path = os.path.join(test_dir, "model.safetensors")
config_data = transformer_test_env["config"].copy()
# case 1: tie weight
config_data = TINY_CONFIG.copy()
config_data["tie_word_embeddings"] = True
config_path = os.path.join(test_dir, "config.json")
@@ -107,7 +73,6 @@ def test_model_save_load_with_tie_weight(transformer_test_env):
assert model.lm_head.weight.data_ptr() == model.embed_tokens.weight.data_ptr()
assert "lm_head.weight" not in model.state_dict()
# case 2: not tie weight (form tie-weight state dict load)
config_data["tie_word_embeddings"] = False
with open(config_path, "w") as f:
json.dump(config_data, f)