Hugging Face's logo

krishnah27
/
AriaLM

Text Generation
TensorBoard
ONNX
PyTorch
music
music-generation
midi
custom
Model card Files
xet
 Metrics Community
AriaLM / docs /LLD.md
krishnah27's picture
krishnah27
Upload folder using huggingface_hub
30e9297 verified
preview code
|
Raw
History Blame Contribute Delete
3.39 kB

Low-Level Design (LLD)

Module Details

s02_tokenizer.py β€” REMI Tokenizer

Vocabulary Layout (485 tokens total):

[0]        PAD
[1]        BOS (Beginning of Sequence)
[2]        EOS (End of Sequence)
[3]        SEP (Separator)
[4-131]    NoteOn (MIDI pitch 0-127)
[132-259]  NoteOff (MIDI pitch 0-127)
[260-291]  Velocity (32 bins, each = 4 MIDI velocity units)
[292-391]  TimeShift (10ms steps, 10ms-1000ms)
[392-451]  Tempo (40-200 BPM, 60 bins)
[452-483]  Position (32 positions per bar)
[484]      Bar (bar delimiter)

Encoding Algorithm:

  1. Collect all non-drum notes across instruments
  2. Sort by onset time, then pitch
  3. For each note: emit [TimeShift, Position, Velocity, NoteOn, TimeShift(duration), NoteOff]
  4. Insert Bar tokens at measure boundaries
  5. Wrap with BOS/EOS

s04_model.py β€” MusicTransformer

Layer Stack (per block):

Input β†’ RMSNorm β†’ GQA (with RoPE) β†’ Residual Add
      β†’ RMSNorm β†’ SwiGLU FFN       β†’ Residual Add β†’ Output

Grouped Query Attention:

  • 8 query heads, 4 key-value heads
  • Each KV head serves 2 query heads (n_rep = 2)
  • Head dim = 256/8 = 32
  • Uses PyTorch 2.0 scaled_dot_product_attention (Flash Attention backend when available)

RoPE Implementation:

  • Precompute sin/cos frequencies: freq[i] = 1 / (ΞΈ^(2i/d))
  • Apply rotation: q' = q * cos + rotate_half(q) * sin
  • Device-compatible real-valued implementation (no complex tensors)

SwiGLU FFN:

  • output = W2(SiLU(W1(x)) * W3(x))
  • Hidden dim = 448 (nearest multiple of 64 to 2/3 * 4 * 256)

s05_trainer.py β€” Training Loop

Optimization Details:

  • AdamW: β₁=0.9, Ξ²β‚‚=0.95, wd=0.1
  • LR schedule: linear warmup (200 steps) β†’ cosine decay to 1e-6
  • Gradient accumulation: 4 steps (effective batch = 4 Γ— 4 = 16)
  • Max gradient norm: 1.0

Memory Budget (estimated for 512 seq len):

Component Memory
Model weights (FP32) ~20MB
Gradients ~20MB
Optimizer states ~40MB
Activations (w/ ckpt) ~50MB
Data batch ~4MB
Total ~134MB

s06_generator.py β€” Inference

Decoding Pipeline:

logits β†’ temperature_scale β†’ repetition_penalty β†’ top_k_filter β†’ top_p_filter β†’ softmax β†’ multinomial_sample

KV-Cache:

  • Each layer stores (K, V) tensors after each forward pass
  • New tokens only compute attention against cached KV + new KV
  • Reset between generations to prevent cross-contamination

Class Diagram 

MusicTokenizer           MusicTransformer              Trainer
β”œβ”€β”€ midi_to_tokens()     β”œβ”€β”€ TransformerBlock (Γ—6)     β”œβ”€β”€ _train_epoch()
β”œβ”€β”€ tokens_to_midi()     β”‚   β”œβ”€β”€ GroupedQueryAttention  β”œβ”€β”€ _validate()
β”œβ”€β”€ decode_token()       β”‚   β”‚   β”œβ”€β”€ wq, wk, wv, wo    β”œβ”€β”€ _save_checkpoint()
β”œβ”€β”€ save()/load()        β”‚   β”‚   β”œβ”€β”€ RoPE application   └── load_checkpoint()
β”‚                        β”‚   β”‚   └── KV-cache
└── Token constants      β”‚   β”œβ”€β”€ SwiGLU FFN
                         β”‚   └── RMSNorm (Γ—2)
                         β”œβ”€β”€ token_emb (weight-tied)
                         β”œβ”€β”€ output projection
                         └── forward() / from_config()