| from __future__ import annotations |
|
|
| import copy |
|
|
| import pytest |
| import torch |
| from torch import nn |
|
|
| from music3lab.codec.external_finetune import ( |
| ExternalSource, |
| FineTuneMetrics, |
| WarmupCosineSchedule, |
| deterministic_guarded_crop, |
| evaluate_gates, |
| load_source_splits, |
| mixed_batch_loss, |
| select_checkpoint, |
| serialize_source_splits, |
| split_external_sources, |
| ) |
|
|
|
|
| USER22 = tuple(f"user22-{index:02d}" for index in range(22)) |
|
|
|
|
| def _sources() -> tuple[ExternalSource, ...]: |
| """Two thousand admissible sources plus user music that must stay OOD.""" |
| return tuple( |
| ExternalSource(source_id=f"external-{index:04d}", samples=44100 * 12) |
| for index in range(2000) |
| ) + tuple( |
| ExternalSource(source_id=source_id, samples=44100 * 12, user_provided=True) |
| for source_id in USER22 |
| ) |
|
|
|
|
| def test_source_exclusive_1600_200_200_split_and_guarded_crops_are_deterministic() -> None: |
| first = split_external_sources(_sources(), seed=73) |
| second = split_external_sources(_sources(), seed=73) |
|
|
| assert {name: len(rows) for name, rows in first.items()} == { |
| "train": 1600, |
| "validation": 200, |
| "heldout": 200, |
| } |
| assert first == second |
| owners = { |
| row.source_id: name for name, rows in first.items() for row in rows |
| } |
| assert len(owners) == 2000 |
| assert not set(USER22) & set(owners) |
|
|
| crop = deterministic_guarded_crop( |
| first["train"][0], split="train", epoch=3, item_index=7, seed=73 |
| ) |
| assert crop == deterministic_guarded_crop( |
| first["train"][0], split="train", epoch=3, item_index=7, seed=73 |
| ) |
| assert crop.samples == 44032 |
| assert 44100 * 5 <= crop.start_sample |
| assert crop.start_sample + crop.samples <= 44100 * 7 |
| assert crop.source_id in owners and owners[crop.source_id] == "train" |
|
|
|
|
| def test_manifest_loader_rejects_user_audio_from_train_or_checkpoint_selection(tmp_path) -> None: |
| manifest = split_external_sources(_sources(), seed=9) |
| path = tmp_path / "sources.json" |
| path.write_text(serialize_source_splits(manifest), encoding="utf-8") |
|
|
| loaded = load_source_splits(path) |
| assert loaded == manifest |
| assert not set(USER22) & { |
| row.source_id for rows in loaded.values() for row in rows |
| } |
| with pytest.raises(ValueError, match="user-provided"): |
| split_external_sources( |
| _sources(), seed=9, train_source_ids=("external-0000", USER22[0]) |
| ) |
|
|
|
|
| def test_mixed_loss_uses_six_external_two_music3_and_exact_weights() -> None: |
| teacher_nmse = torch.tensor(2.0) |
| teacher_ruler = torch.tensor(3.0) |
| external_ruler = torch.tensor(5.0) |
| prior = torch.tensor(7.0) |
| result = mixed_batch_loss( |
| teacher_nmse=teacher_nmse, |
| teacher_ruler=teacher_ruler, |
| external_ruler=external_ruler, |
| external_prior=prior, |
| external_count=6, |
| music3_count=2, |
| ) |
|
|
| assert result.external_count == 6 |
| assert result.music3_count == 2 |
| assert float(result.total) == pytest.approx( |
| 0.70 * 5.0 + 0.30 * (2.0 + 0.05 * 3.0) + 0.005 * 7.0 |
| ) |
| assert result.teacher_weight == pytest.approx(0.30) |
| assert result.external_weight == pytest.approx(0.70) |
|
|
|
|
| def test_only_encoder_receives_gradients_and_frozen_decoder_is_unchanged() -> None: |
| encoder = nn.Linear(2, 2, bias=False) |
| decoder = nn.Linear(2, 2, bias=False) |
| for parameter in decoder.parameters(): |
| parameter.requires_grad_(False) |
| decoder_before = copy.deepcopy(decoder.state_dict()) |
|
|
| teacher_audio = torch.tensor([[1.0, -1.0], [0.5, 0.25]]) |
| external_audio = torch.tensor([[0.25, -0.5]]).repeat(6, 1) |
| predicted = encoder(torch.cat((teacher_audio, external_audio), dim=0)) |
| rendered = decoder(predicted) |
| result = mixed_batch_loss( |
| teacher_nmse=(predicted[:2] - teacher_audio).square().mean(), |
| teacher_ruler=(rendered[:2] - teacher_audio).square().mean(), |
| external_ruler=(rendered[2:] - external_audio).square().mean(), |
| external_prior=predicted[2:].square().mean(), |
| external_count=6, |
| music3_count=2, |
| ) |
| result.total.backward() |
|
|
| assert all(parameter.grad is not None for parameter in encoder.parameters()) |
| assert all(parameter.grad is None for parameter in decoder.parameters()) |
| assert decoder.state_dict().keys() == decoder_before.keys() |
| assert all(torch.equal(value, decoder_before[key]) for key, value in decoder.state_dict().items()) |
|
|
|
|
| def test_warmup_cosine_endpoints_are_frozen() -> None: |
| schedule = WarmupCosineSchedule( |
| warmup_steps=10, |
| total_steps=100, |
| maximum_learning_rate=1e-3, |
| minimum_learning_rate=1e-5, |
| ) |
| assert schedule(0) == pytest.approx(0.0) |
| assert schedule(10) == pytest.approx(1e-3) |
| assert schedule(100) == pytest.approx(1e-5) |
| assert schedule(50) < schedule(10) |
|
|
|
|
| def test_synthetic_training_improves_external_ruler_without_teacher_regression() -> None: |
| scale = nn.Parameter(torch.zeros(())) |
| optimizer = torch.optim.SGD([scale], lr=0.15) |
| baseline = FineTuneMetrics(teacher_ruler=1.0, external_ruler=1.1025) |
| for _ in range(50): |
| optimizer.zero_grad() |
| loss = mixed_batch_loss( |
| teacher_nmse=(scale - 1.0).square(), |
| teacher_ruler=(scale - 1.0).square(), |
| external_ruler=(scale - 1.05).square(), |
| external_prior=scale.square(), |
| external_count=6, |
| music3_count=2, |
| ) |
| loss.total.backward() |
| optimizer.step() |
| candidate = FineTuneMetrics( |
| teacher_ruler=float((scale.detach() - 1.0).square()), |
| external_ruler=float((scale.detach() - 1.05).square()), |
| ) |
| gate = evaluate_gates(baseline=baseline, candidate=candidate) |
|
|
| assert gate.external_ruler_improved is True |
| assert gate.teacher_regression_fraction <= 0.05 |
| assert gate.teacher_regression_within_limit is True |
| assert gate.passes is True |
|
|
|
|
| def test_heldout_gates_and_user22_ood_report_cannot_select_a_checkpoint() -> None: |
| baseline = FineTuneMetrics(teacher_ruler=1.0, external_ruler=4.0) |
| checkpoints = { |
| "step-010": FineTuneMetrics(teacher_ruler=1.02, external_ruler=2.0), |
| "step-020": FineTuneMetrics(teacher_ruler=1.01, external_ruler=2.1), |
| } |
| selected = select_checkpoint( |
| baseline=baseline, |
| validation=checkpoints, |
| heldout={"step-010": FineTuneMetrics(teacher_ruler=4.0, external_ruler=0.1)}, |
| user22_ood={"step-010": FineTuneMetrics(teacher_ruler=0.0, external_ruler=0.0)}, |
| ) |
|
|
| assert selected.name == "step-010" |
| assert selected.selection_split == "validation" |
| assert selected.heldout_gate is not None |
| assert selected.user22_ood_report is not None |
| assert selected.user22_ood_report.influenced_selection is False |
|
|
