Skip to content

feat: dual vocab embeddings#3

Merged
imabeastdrew merged 2 commits intomainfrom
dual-vocab
Dec 3, 2025
Merged

feat: dual vocab embeddings#3
imabeastdrew merged 2 commits intomainfrom
dual-vocab

Conversation

@schennam714
Copy link
Copy Markdown
Collaborator


Summary

  • Separate embedding tables for melody and chord tokens - Models now use dedicated melody_embed and chord_embed layers instead of a single unified embedding, allowing each modality to develop task-appropriate representations
  • Removed runtime logits masking - Output layers now project directly to chord_vocab_size, making chord-only output an architectural constraint rather than a runtime mask
  • Fixed eval decoding bug - Chord predictions now correctly decode using chord_id_to_token instead of unified_id_to_token
  • Preprocessing now saves separate vocab files - vocab_melody.json and vocab_chord.json are saved alongside vocab_unified.json

Changes

Models (models/offline.py, models/online.py):

  • Offline: Custom encoder_embed for melody input, T5's internal embedding for decoder (chord) tokens
  • Online: melody_embed + chord_embed selected via is_melody boolean mask
  • Both models output logits over chord_vocab_size only (no masking needed)
  • Removed _ChordOnlyLogitsProcessor and logits_mask buffer

Data (data/base.py, data/online.py):

  • BaseDataset now loads separate vocab_melody.json and vocab_chord.json
  • Added _unified_to_chord_id() and _unified_to_melody_id() conversion methods
  • OnlineDataset uses base class vocab sizes (removed redundant recalculation)

Eval (eval/offline.py, eval/online.py):

  • Fixed vocab decoding: chords use chord_id_to_token, melody uses melody_id_to_token

Preprocessing (scripts/preprocess.py):

  • Added melody_vocab.save() and chord_vocab.save() calls

Tests:

  • Added semantic validation tests ensuring chords decode as chord tokens
    (not melody)
  • Updated all model tests for new dual-vocab API

Test plan

  • All 54 unit/integration tests pass
  • Ruff linting passes
  • mypy type checking passes
  • Train offline model and verify convergence
  • Train online model and verify convergence
  • Run eval on trained checkpoints

@imabeastdrew imabeastdrew merged commit 85e72e1 into main Dec 3, 2025
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants