File size: 4,798 Bytes
90884df | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | from __future__ import annotations
import math
import pytest
import torch
from music3lab.codec.native_state_multiclip import (
CAPABILITY_LABEL,
LiteNativeStateBatch,
ScaleUpMetrics,
assert_no_residual_loss_path,
deterministic_split_plan,
derive_feedback_deterministically,
extract_global_from_fused,
scale_up_gate,
validate_lite_native_state_batch,
)
FRAMES = 25
WIDTH = 4096
def _batch(n: int = 2) -> LiteNativeStateBatch:
global_hidden = torch.randn(n, FRAMES, WIDTH).bfloat16()
return LiteNativeStateBatch(
audio=torch.randn(n, 2, 44_032, dtype=torch.float32).clamp(-1, 1),
continuous_latent=torch.randn(n, 128, 86).bfloat16(),
semantic=torch.randint(0, 16_384, (n, FRAMES), dtype=torch.int64),
residual=torch.randint(0, 2048, (n, FRAMES, 7), dtype=torch.int64),
priming=torch.randint(0, 16_384, (n, 8), dtype=torch.int64),
global_hidden=global_hidden,
feedback=torch.randn(n, FRAMES, WIDTH).bfloat16(),
seeds=torch.arange(n, dtype=torch.int64),
)
def test_lite_schema_accepts_exact_tensors_and_rejects_bad_dtype_shape_nonfinite() -> None:
batch = _batch()
validate_lite_native_state_batch(batch)
with pytest.raises(ValueError, match="audio.*float32"):
validate_lite_native_state_batch(LiteNativeStateBatch(**{**batch.__dict__, "audio": batch.audio.bfloat16()}))
with pytest.raises(ValueError, match="continuous_latent.*128, 86"):
validate_lite_native_state_batch(LiteNativeStateBatch(**{**batch.__dict__, "continuous_latent": batch.continuous_latent[..., :-1]}))
bad_feedback = batch.feedback.clone()
bad_feedback[0, 0, 0] = float("nan")
with pytest.raises(ValueError, match="feedback.*finite"):
validate_lite_native_state_batch(LiteNativeStateBatch(**{**batch.__dict__, "feedback": bad_feedback}))
def test_split_plan_is_deterministic_disjoint_and_has_the_bounded_composition() -> None:
first = deterministic_split_plan()
assert first == deterministic_split_plan()
assert len(first.train) == 80
assert len(first.validation) == len(first.heldout) == 12
assert len(first.diagnostic_train) == len(first.diagnostic_heldout) == 16
assert sum(item.source == "reused_french_house" for item in first.train) == 64
assert sum(item.source == "fresh" for item in first.train) == 16
assert len({item.prompt_id for item in first.train}) == 2
primary = first.train + first.validation + first.heldout
assert not ({item.prompt_id for item in first.validation} & {item.prompt_id for item in first.train})
assert not ({item.prompt_id for item in first.heldout} & {item.prompt_id for item in first.train})
assert not ({item.prompt_id for item in first.validation} & {item.prompt_id for item in first.heldout})
assert len({(item.prompt_id, item.seed) for item in primary}) == len(primary)
assert all(item.frames == FRAMES for item in primary)
def test_global_extraction_is_exact_fused_prefix_and_feedback_callback_consumes_no_rng() -> None:
batch = _batch()
fused = torch.randn(batch.audio.shape[0], FRAMES, 32_768).bfloat16()
fused[..., :WIDTH] = batch.global_hidden
assert torch.equal(extract_global_from_fused(fused), batch.global_hidden)
state_before = torch.random.get_rng_state()
observed = []
def callback(global_hidden: torch.Tensor) -> torch.Tensor:
observed.append(global_hidden.shape)
return global_hidden + 1
feedback = derive_feedback_deterministically(batch.global_hidden, callback)
assert observed == [batch.global_hidden.shape]
assert feedback.shape == (2, FRAMES, WIDTH)
assert torch.equal(feedback, batch.global_hidden + 1)
assert torch.equal(state_before, torch.random.get_rng_state())
def test_residual_is_provenance_only_not_a_loss_target() -> None:
targets = assert_no_residual_loss_path(_batch())
assert set(targets) == {"semantic", "feedback", "global_hidden"}
assert all("residual" not in name for name in targets)
def test_scaleup_gate_requires_all_generalization_and_capability_conditions() -> None:
baseline = ScaleUpMetrics(1.0, 1.0, 1.0, 1.0, math.log(16_384))
passing = ScaleUpMetrics(0.80, 0.94, 0.80, 0.94, math.log(16_384) - 0.11)
assert scale_up_gate(baseline, passing).accepted is True
bad_gap = ScaleUpMetrics(0.40, 0.94, 0.40, 0.94, math.log(16_384) - 0.11)
result = scale_up_gate(baseline, bad_gap)
assert result.accepted is False
assert "train_vs_heldout_gap" in result.reasons
def test_capability_is_truthfully_labeled_as_captured_multiclip_only() -> None:
assert CAPABILITY_LABEL == "captured_music3_multiclip_only"
assert CAPABILITY_LABEL != "native_tokenizer"
assert CAPABILITY_LABEL != "arbitrary_external_music"
|