| import numpy as np |
| import pytest |
| import torch |
|
|
| import _net |
| import ssl_mamba as S |
|
|
|
|
| def test_pad_long_truncated(): |
| x = np.arange(100000, dtype=np.float32) |
| assert np.array_equal(_net.pad_repeat(x), x[:64600]) |
|
|
|
|
| def test_pad_short_repeated(): |
| x = np.arange(1000, dtype=np.float32) |
| out = _net.pad_repeat(x) |
| assert out.shape == (64600,) and np.array_equal(out[:1000], x) and np.array_equal(out[1000:2000], x) |
|
|
|
|
| @pytest.fixture(scope="module") |
| def model(): |
| m = S.SSLMamba(); m.load() |
| yield m |
| m.unload() |
|
|
|
|
| def test_eval_mode_and_params(model): |
| assert not any(mod.training for mod in model.net.modules()) |
| assert sum(p.numel() for p in model.net.parameters()) == 317_521_282 |
| assert sum(t.numel() for t in model.net.state_dict().values()) == 317_521_282 |
|
|
|
|
| def test_batch_matches_single(model): |
| rng = np.random.default_rng(0) |
| a = [rng.standard_normal(n).astype(np.float32) * s for n, s in ((48000, 0.1), (20000, 0.5), (70000, 0.02), (3000, 0.3))] |
| batched = model.score_batch(a, [16000] * 4) |
| single = [model.score_batch([x], [16000])[0] for x in a] |
| assert len(batched) == 4 and all(np.isfinite(batched)) |
| |
| assert np.allclose(batched, single, rtol=1e-4, atol=1e-4) |
|
|
|
|
| def test_chunked_scan_matches_reference(model): |
| rng = np.random.default_rng(1) |
| a = [rng.standard_normal(50000).astype(np.float32) * 0.1 for _ in range(3)] |
| s1 = model.score_batch(a, [16000] * 3) |
| model.net.set_ref_scan(True) |
| try: |
| s2 = model.score_batch(a, [16000] * 3) |
| finally: |
| model.net.set_ref_scan(False) |
| assert np.allclose(s1, s2, atol=1e-4) |
|
|
|
|
| def test_deterministic(model): |
| a = [np.random.randn(40000).astype(np.float32) for _ in range(2)] |
| assert model.score_batch(a, [16000] * 2) == model.score_batch(a, [16000] * 2) |
|
|
|
|
| def test_sign_and_finite(model): |
| |
| a = [np.random.default_rng(2).standard_normal(64600).astype(np.float32) * 0.05] |
| s = model.score_batch(a, [16000])[0] |
| x = torch.from_numpy(_net.pad_repeat(a[0])[None]).to(model.device) |
| with torch.no_grad(): |
| out = model.net(x)[0] |
| assert np.isfinite(s) and abs(s - float(out[1] - out[0])) < 1e-4 |
|
|
