Add tests for MusicBERT token-to-note pooling and REMI↔BPE alignment - #6
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds comprehensive test coverage for MusicBERT token-to-note pooling and REMI↔BPE alignment functionality, while also integrating MusicBERT as an optional note encoder into the analysis pipeline. The tests validate numerical behavior of the pooler and API contracts of the alignment builder.
Key changes:
- New test suites for
TokenToNotePoolerandbuild_alignmentfunctions with unit tests covering weight invariance, weighted summation, and error handling - Complete MusicBERT integration including backbone, note encoder, and pooling modules with LoRA adapter support
- Data pipeline enhancement to load and attach alignment data from npz files during graph transformation
Reviewed changes
Copilot reviewed 12 out of 13 changed files in this pull request and generated 13 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_musicbert_pooler.py | Unit tests for token-to-note pooling with weight invariance and summation checks |
| tests/test_remi_bpe_aligner.py | Unit tests for alignment builder and npz loader functionality |
| analysisgnn/modules/token_to_note_pooler.py | New pooler module that aggregates token embeddings into note embeddings using weighted edges |
| analysisgnn/modules/init.py | Module exports for TokenToNotePooler |
| analysisgnn/models/musicbert_backbone.py | MusicBERT backbone wrapper with LoRA adapter support |
| analysisgnn/models/musicbert_note_encoder.py | Note encoder combining MusicBERT backbone with token-to-note pooler |
| analysisgnn/models/init.py | Added exports for MusicBERT components |
| analysisgnn/data/remi_bpe_aligner.py | Alignment builder and loader for REMI↔BPE token-to-note mappings |
| analysisgnn/data/datamodules/analysis.py | Transform pipeline updated to load and attach alignment data from npz files |
| analysisgnn/models/analysis.py | ContinualAnalysisGNN updated to support optional MusicBERT note encoder |
| analysisgnn/train/train_analysisgnn.py | CLI arguments and training logic for MusicBERT configuration |
| requirements.txt | Added transformers, peft, and miditok dependencies |
| environment.yml | Added transformers, peft, and miditok dependencies |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| class MusicBertNoteEncoder(nn.Module): | ||
| def __init__( | ||
| self, | ||
| pretrained_name: str, | ||
| adapter_cfg: Optional[MusicBertAdapterConfig] = None, | ||
| freeze_backbone: bool = True, | ||
| ) -> None: |
There was a problem hiding this comment.
The MusicBertNoteEncoder class lacks documentation explaining its purpose and how it combines the backbone and pooler. Consider adding a docstring that describes the two-stage process: token embedding via MusicBERT followed by token-to-note pooling.
There was a problem hiding this comment.
@copilot open a new pull request to apply changes based on this feedback
| class TokenToNotePooler(nn.Module): | ||
| def forward( | ||
| self, | ||
| token_states: torch.Tensor, | ||
| token2note: List[torch.Tensor], | ||
| num_notes: List[int], | ||
| ) -> Tuple[torch.Tensor, torch.BoolTensor]: |
There was a problem hiding this comment.
The TokenToNotePooler class lacks a docstring explaining its purpose, parameters, and return values. Adding documentation would help users understand the pooling mechanism and expected input/output formats, especially the edge weight format [token_idx, note_idx, weight].
There was a problem hiding this comment.
@copilot open a new pull request to apply changes based on this feedback
documentation for bpe alignment function. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
|
@manoskary I've opened a new pull request, #7, to work on those changes. Once the pull request is ready, I'll request review from you. |
fix inconsistent flag logic. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
|
@manoskary I've opened a new pull request, #8, to work on those changes. Once the pull request is ready, I'll request review from you. |
Co-authored-by: manoskary <26930454+manoskary@users.noreply.github.com>
Add documentation to MusicBertNoteEncoder class
Co-authored-by: manoskary <26930454+manoskary@users.noreply.github.com>
Add docstring to TokenToNotePooler class
Motivation
Description
tests/test_musicbert_pooler.pywith checks thatTokenToNotePooleris invariant to edge weight splitting and correctly sums weighted token contributions.tests/test_remi_bpe_aligner.pywhich verifiesbuild_alignmentraisesNotImplementedErrorwithout mappings and returns the expectedBpeNoteAlignmentfields when provided a mapping.TokenToNotePoolerandbuild_alignmentimplementations.Testing
pytest -qto collect and run tests in the repository.ModuleNotFoundErrorfortorchandnumpy, so the new tests were not executed; this indicates required packages are missing from the test environment.tests/test_musicbert_pooler.pyandtests/test_remi_bpe_aligner.py.torchandnumpyare available in the environment and re-runpytest -q.Codex Task