refactor : 压缩测试代码，消除重复

- fixture 替代重复实例化和 tokenizer 落盘
- parametrize 合并同构测试
- helper 消除 save_h5 + DatasetFactory.load 样板
- 净减 272 行

tests/module/test_encoder.py

@@ -1,6 +1,13 @@
+import json
+import os
+import tempfile
+
+import pytest
+import safetensors.torch as st
 import torch
 
 from astrai.config.model_config import EncoderConfig
+from astrai.model.automodel import AutoModel
 from astrai.model.encoder import EmbeddingEncoder
 
 TINY_CONFIG = dict(
@@ -14,92 +21,56 @@ TINY_CONFIG = dict(
     norm_eps=1e-5,
 )
 
+_device = "cuda" if torch.cuda.is_available() else "cpu"
 
-def test_encoder_forward_mean():
-    config = EncoderConfig(**TINY_CONFIG)
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = EmbeddingEncoder(config).to(device=device)
+
+def _make_model(**kwargs):
+    config = EncoderConfig(**{**TINY_CONFIG, **kwargs})
+    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)
     model.eval()
 
     batch_size, seq_len = 2, 8
     input_ids = torch.randint(
-        0, config.vocab_size, (batch_size, seq_len), device=device
+        0, TINY_CONFIG["vocab_size"], (batch_size, seq_len), device=_device
     )
 
     with torch.no_grad():
         output = model(input_ids)
 
-    assert output.shape == (batch_size, config.dim)
-    assert not torch.isnan(output).any()
-
-
-def test_encoder_forward_cls():
-    config = EncoderConfig(**{**TINY_CONFIG, "pooling_type": "cls"})
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = EmbeddingEncoder(config).to(device=device)
-    model.eval()
-
-    batch_size, seq_len = 2, 8
-    input_ids = torch.randint(
-        0, config.vocab_size, (batch_size, seq_len), device=device
-    )
-
-    with torch.no_grad():
-        output = model(input_ids)
-
-    assert output.shape == (batch_size, config.dim)
-    assert not torch.isnan(output).any()
-
-
-def test_encoder_forward_last():
-    config = EncoderConfig(**{**TINY_CONFIG, "pooling_type": "last"})
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = EmbeddingEncoder(config).to(device=device)
-    model.eval()
-
-    batch_size, seq_len = 2, 8
-    input_ids = torch.randint(
-        0, config.vocab_size, (batch_size, seq_len), device=device
-    )
-
-    with torch.no_grad():
-        output = model(input_ids)
-
-    assert output.shape == (batch_size, config.dim)
+    assert output.shape == (batch_size, TINY_CONFIG["dim"])
     assert not torch.isnan(output).any()
 
 
 def test_encoder_forward_with_padding():
-    config = EncoderConfig(**TINY_CONFIG)
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = EmbeddingEncoder(config).to(device=device)
+    model = _make_model()
     model.eval()
 
     batch_size, seq_len = 2, 8
     input_ids = torch.randint(
-        0, 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():
         output = model(input_ids, input_mask=input_mask)
 
-    assert output.shape == (batch_size, config.dim)
+    assert output.shape == (batch_size, TINY_CONFIG["dim"])
     assert not torch.isnan(output).any()
 
 
 def test_encoder_normalize():
-    config = EncoderConfig(
-        **{**TINY_CONFIG, "pooling_type": "mean", "normalize_embeddings": True}
-    )
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = EmbeddingEncoder(config).to(device=device)
+    model = _make_model(pooling_type="mean", normalize_embeddings=True)
     model.eval()
 
     batch_size, seq_len = 2, 8
     input_ids = torch.randint(
-        0, config.vocab_size, (batch_size, seq_len), device=device
+        0, TINY_CONFIG["vocab_size"], (batch_size, seq_len), device=_device
     )
 
     with torch.no_grad():
@@ -110,24 +81,19 @@ def test_encoder_normalize():
 
 
 def test_encoder_register():
-    from astrai.model.automodel import AutoModel
-
     assert AutoModel.is_registered("embedding")
     cls = AutoModel.get_component_class("embedding")
     assert cls is EmbeddingEncoder
 
 
 def test_encoder_from_transformer_checkpoint():
-    config = EncoderConfig(**TINY_CONFIG)
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = EmbeddingEncoder(config).to(device=device)
-
+    model = _make_model()
     state_dict = model.state_dict()
     state_dict["lm_head.weight"] = torch.randn(
-        config.vocab_size, config.dim, device=device
+        TINY_CONFIG["vocab_size"], TINY_CONFIG["dim"], device=_device
     )
 
-    new_model = EmbeddingEncoder(config).to(device=device)
+    new_model = _make_model()
     new_model.load_state_dict(state_dict, strict=True)
 
     for key in model.state_dict():
@@ -135,12 +101,6 @@ def test_encoder_from_transformer_checkpoint():
 
 
 def test_encoder_save_load():
-    import json
-    import os
-    import tempfile
-
-    import safetensors.torch as st
-
     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")