mirror of
https://github.com/RYDE-WORK/lnp_ml.git
synced 2026-09-18 21:13:20 +08:00
105 lines
3.3 KiB
Python
105 lines
3.3 KiB
Python
"""ResidualConcatFusion 模块自检脚本
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import sys
|
||
import traceback
|
||
from typing import Callable, List, Tuple
|
||
|
||
import torch
|
||
|
||
from lnp_ml.modeling.layers.fusion import ResidualConcatFusion
|
||
|
||
|
||
def test_full_tokens_shape() -> None:
|
||
fusion = ResidualConcatFusion(d_model=256)
|
||
chem = torch.randn(2, 4, 256)
|
||
tab = torch.randn(2, 4, 256)
|
||
f_moe = torch.randn(2, 256)
|
||
f_llm = torch.randn(2, 256)
|
||
out = fusion(chem, tab, f_moe, f_llm)
|
||
assert out.shape == (2, 256), f"out shape {out.shape} != (2, 256)"
|
||
assert fusion.fusion_dim == 256
|
||
|
||
|
||
def test_without_extras() -> None:
|
||
"""关闭 MoE 和 LLM 时只拼接 chem + tab。"""
|
||
fusion = ResidualConcatFusion(d_model=64)
|
||
chem = torch.randn(3, 4, 64)
|
||
tab = torch.randn(3, 4, 64)
|
||
assert fusion(chem, tab).shape == (3, 64)
|
||
|
||
|
||
def test_nchem3_with_moe() -> None:
|
||
fusion = ResidualConcatFusion(d_model=64)
|
||
chem = torch.randn(2, 3, 64)
|
||
tab = torch.randn(2, 4, 64)
|
||
f_moe = torch.randn(2, 64)
|
||
assert fusion(chem, tab, f_moe=f_moe).shape == (2, 64)
|
||
|
||
|
||
def test_attn_weights_shape() -> None:
|
||
"""注意力只池化真实 token：序列长度 = 4 + 4 = 8；f_moe/f_llm 走门控残差，不参与 softmax。"""
|
||
fusion = ResidualConcatFusion(d_model=64)
|
||
chem = torch.randn(2, 4, 64)
|
||
tab = torch.randn(2, 4, 64)
|
||
f_moe = torch.randn(2, 64)
|
||
f_llm = torch.randn(2, 64)
|
||
out, weights = fusion(chem, tab, f_moe, f_llm, return_attn_weights=True)
|
||
assert out.shape == (2, 64)
|
||
assert weights.shape == (2, 8), f"weights shape {weights.shape} != (2, 8)"
|
||
|
||
|
||
def test_concat_strategy_rejected() -> None:
|
||
try:
|
||
ResidualConcatFusion(d_model=64, strategy="concat")
|
||
except ValueError:
|
||
return
|
||
raise AssertionError("concat 策略应被拒绝")
|
||
|
||
|
||
def test_grad_flows() -> None:
|
||
fusion = ResidualConcatFusion(d_model=32)
|
||
chem = torch.randn(2, 4, 32, requires_grad=True)
|
||
tab = torch.randn(2, 4, 32, requires_grad=True)
|
||
fusion(chem, tab).sum().backward()
|
||
assert chem.grad is not None and torch.isfinite(chem.grad).all()
|
||
|
||
|
||
# ──────────────────────────────────────────────────────────────────────
|
||
# Runner
|
||
# ──────────────────────────────────────────────────────────────────────
|
||
|
||
def _collect_tests() -> List[Tuple[str, Callable[[], None]]]:
|
||
return [
|
||
(name, fn) for name, fn in globals().items()
|
||
if name.startswith("test_") and callable(fn)
|
||
]
|
||
|
||
|
||
def main() -> int:
|
||
tests = _collect_tests()
|
||
pass_count = 0
|
||
fail: List[Tuple[str, str]] = []
|
||
for name, fn in tests:
|
||
try:
|
||
fn()
|
||
except Exception:
|
||
fail.append((name, traceback.format_exc()))
|
||
print(f"[FAIL] {name}")
|
||
else:
|
||
pass_count += 1
|
||
print(f"[ OK ] {name}")
|
||
print(f"\n{'='*60}")
|
||
print(f"Passed: {pass_count}/{len(tests)}")
|
||
if fail:
|
||
print(f"\n{'='*60}\nFailures:")
|
||
for name, tb in fail:
|
||
print(f"\n--- {name} ---\n{tb}")
|
||
return 1
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main()) |