refactor: eliminate test duplication via shared helpers

- 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

tests/module/test_encoder.py

@@ -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())