Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,24 @@
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+*.egg-info/
+dist/
+build/
+.eggs/
+*.egg
+.env
+.venv
+env/
+venv/
+data/
+checkpoints/
+output/
+*.mid
+*.midi
+runs/
+*.pt
+*.pth
+*.ckpt
+.DS_Store
+*.log

Dockerfile

@@ -0,0 +1,13 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+RUN mkdir -p data checkpoints output runs
+
+ENTRYPOINT ["python", "-m", "src.s00_main"]
+CMD ["train+generate"]

README.md

@@ -1,3 +1,98 @@
----
-license: mit
----
+# Music Generation LLM
+
+A LLaMA-style Transformer model for symbolic music generation, trained on MIDI data.
+
+## Architecture
+
+- **Model**: LLaMA-style Transformer with RoPE, GQA, SwiGLU, RMSNorm
+- **Tokenizer**: REMI (REvamped MIDI-derived) — SOTA for symbolic music
+- **Dataset**: `drengskapur/midi-classical-music` — 4,796 classical MIDI files (~50MB)
+- **Training**: AdamW + Cosine LR warmup + AMP + Gradient Checkpointing
+
+## Key Features
+
+- **Memory Efficient**: Grouped Query Attention, gradient checkpointing, mixed precision
+- **OOM Safe**: Conservative batch sizes, AMP, lazy data loading
+- **SOTA Techniques**: RoPE, SwiGLU, RMSNorm, KV-cache, top-p/top-k sampling
+
+## Quick Start
+
+```bash
+# Install dependencies
+pip install -r requirements.txt
+
+# Train + Generate (default)
+python3 -m src.s00_main train+generate
+
+# Train only
+python3 -m src.s00_main train --epochs 20 --batch-size 4
+
+# Generate from checkpoint
+python3 -m src.s00_main generate --temperature 0.85
+```
+
+## Project Structure
+
+```
+music_gen_llm/
+├── src/
+│   ├── s00_main.py          # Entry point — orchestrates pipeline
+│   ├── s01_config.py         # All configuration dataclasses
+│   ├── s02_tokenizer.py      # REMI MIDI tokenizer
+│   ├── s03_dataset.py        # Data download + tokenization + DataLoader
+│   ├── s04_model.py          # MusicTransformer (LLaMA-style)
+│   ├── s05_trainer.py        # Training loop with AMP + checkpointing
+│   ├── s06_generator.py      # Autoregressive generation with KV-cache
+│   └── s07_utils.py          # Logging, memory monitoring, seeding
+├── tests/
+│   └── test_pipeline.py      # 7 unit tests covering all components
+├── scripts/
+│   ├── download_data.sh      # Dataset setup
+│   ├── train.sh              # Training launcher
+│   └── generate.sh           # Generation launcher
+├── docs/
+│   ├── README.md             # This file
+│   ├── HLD.md                # High-Level Design
+│   ├── LLD.md                # Low-Level Design
+│   └── flow_diagram.drawio   # Execution flow diagram
+├── data/                     # Downloaded MIDI + tokenized cache
+├── checkpoints/              # Saved model weights
+├── output/                   # Generated MIDI files
+├── requirements.txt
+├── Dockerfile
+└── .gitignore
+```
+
+## Execution Flow
+
+```
+s00_main.py → s01_config.py → s02_tokenizer.py → s03_dataset.py → s04_model.py → s05_trainer.py → s06_generator.py
+     │              │                │                  │                │               │                │
+  Entry point   Load configs   Init tokenizer   Download & tokenize  Build model  Train loop     Generate MIDI
+```
+
+## Model Specifications
+
+| Parameter     | Value                              |
+|---------------|-------------------------------------|
+| Dim           | 256                                |
+| Layers        | 6                                  |
+| Heads         | 8 (Q) / 4 (KV) — GQA              |
+| Hidden (FFN)  | 448 (SwiGLU)                       |
+| Max Seq Len   | 1024                               |
+| Vocab Size    | 485 (REMI tokens)                  |
+| Parameters    | ~5M                                |
+| Precision     | BF16/FP16 (AMP)                    |
+
+## Algorithms Used
+
+1. **Rotary Position Embeddings (RoPE)** — Su et al. 2021
+2. **Grouped Query Attention (GQA)** — Ainslie et al. 2023, from LLaMA-2
+3. **SwiGLU Activation** — Shazeer 2020, from LLaMA
+4. **RMS Layer Normalization** — Zhang & Sennrich 2019
+5. **REMI Tokenization** — Huang & Yang 2020
+6. **Cosine Annealing with Warmup** — Loshchilov & Hutter 2017
+7. **Gradient Checkpointing** — Chen et al. 2016
+8. **KV-Cache** — Standard for efficient autoregressive decoding
+9. **Nucleus (Top-p) + Top-k Sampling** — Holtzman et al. 2020
+10. **Repetition Penalty** — Keskar et al. 2019

data/midi_files/.gitattributes

@@ -0,0 +1,55 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.lz4 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+# Audio files - uncompressed
+*.pcm filter=lfs diff=lfs merge=lfs -text
+*.sam filter=lfs diff=lfs merge=lfs -text
+*.raw filter=lfs diff=lfs merge=lfs -text
+# Audio files - compressed
+*.aac filter=lfs diff=lfs merge=lfs -text
+*.flac filter=lfs diff=lfs merge=lfs -text
+*.mp3 filter=lfs diff=lfs merge=lfs -text
+*.ogg filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text
+# Image files - uncompressed
+*.bmp filter=lfs diff=lfs merge=lfs -text
+*.gif filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.tiff filter=lfs diff=lfs merge=lfs -text
+# Image files - compressed
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.webp filter=lfs diff=lfs merge=lfs -text

data/midi_files/README.md

@@ -0,0 +1,27 @@
+---
+pretty_name: MIDI Classical Music
+dataset_name: midi-classical-music
+tags:
+- music
+- classical
+- midi
+- dataset
+- composers
+- music-analysis
+- music-generation
+license: mit
+language:
+- en
+size_categories:
+- 1K<n<10K
+---
+
+# MIDI Classical Music
+
+This dataset contains a comprehensive collection of MIDI files representing classical music compositions from various renowned composers. 
+
+The collection includes works from composers such as Bach, Beethoven, Chopin, Mozart, and many others. 
+
+The dataset is organized into directories by composer, with each directory containing MIDI files of their compositions.
+
+The dataset is ideal for music analysis, machine learning models for music generation, and other music-related research and applications.

data/tokenized_cache.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b744a77aa1a6d3d25a24f5d1ff4dcfb0071295a45c7b92961fa45469dbb9936
+size 6125102

data/tokenizer.json

@@ -0,0 +1 @@
+{"vocab_size": 485}

docs/HLD.md

@@ -0,0 +1,76 @@
+# High-Level Design (HLD)
+
+## System Overview
+
+Music Generation LLM is a symbolic music generation system that learns patterns from
+classical MIDI music and generates new compositions. It operates entirely on symbolic
+note representations (MIDI events), not raw audio waveforms.
+
+## Architecture Diagram
+
+```
+┌─────────────────────────────────────────────────────────┐
+│                    USER INTERFACE                        │
+│              CLI: train / generate / both                │
+└────────────────────────┬────────────────────────────────┘
+                         │
+┌────────────────────────▼────────────────────────────────┐
+│                  ORCHESTRATOR (s00_main)                 │
+│              Parses args, wires components               │
+└──────┬─────────────────┬────────────────────┬───────────┘
+       │                 │                    │
+┌──────▼──────┐  ┌───────▼───────┐  ┌────────▼────────┐
+│   DATA      │  │    MODEL      │  │   GENERATION    │
+│  PIPELINE   │  │   PIPELINE    │  │   PIPELINE      │
+├─────────────┤  ├───────────────┤  ├─────────────────┤
+│ HF Download │  │ MusicTrans-   │  │ Autoregressive  │
+│ MIDI Parse  │  │ former Build  │  │ w/ KV-cache     │
+│ REMI Token  │  │ Train Loop    │  │ Top-p/k Sample  │
+│ DataLoader  │  │ AMP + Ckpt    │  │ MIDI Export     │
+└─────────────┘  └───────────────┘  └─────────────────┘
+```
+
+## Component Responsibilities
+
+### Data Pipeline (s02 + s03)
+- Download 4,796 MIDI files from HuggingFace
+- Parse with `pretty_midi` library
+- Tokenize using REMI scheme (Note On/Off, Velocity, TimeShift, Bar, Position)
+- Create PyTorch DataLoaders with padding and random cropping
+
+### Model (s04)
+- LLaMA-style Transformer (6 layers, 256 dim, 8 heads)
+- Grouped Query Attention with 4 KV heads (50% memory reduction)
+- SwiGLU feed-forward with RMSNorm
+- RoPE for positional encoding
+- Gradient checkpointing support
+
+### Training (s05)
+- AdamW optimizer with cosine LR warmup
+- Mixed precision (BF16/FP16) via PyTorch AMP
+- Gradient accumulation (effective batch = 32)
+- Early stopping with patience
+- TensorBoard logging
+
+### Generation (s06)
+- Autoregressive decoding with KV-cache
+- Temperature scaling + Top-k + Top-p (nucleus) sampling
+- Repetition penalty for diverse output
+- Direct MIDI file export
+
+## Data Flow
+
+```
+MIDI Files → REMI Tokens → Training → Trained Model → Generation → MIDI Output
+```
+
+## Constraints & Decisions
+
+| Decision                     | Rationale                                    |
+|------------------------------|----------------------------------------------|
+| Symbolic (MIDI) not audio    | 100x smaller data, trainable on consumer GPU |
+| GQA over standard MHA        | 50% KV-cache memory reduction                |
+| Gradient checkpointing       | ~50% memory savings, ~20% speed cost         |
+| BF16 mixed precision         | 50% memory, maintains numeric stability      |
+| REMI tokenization            | Best published results for symbolic music     |
+| 5M params (not 100M+)        | Fits in <=4GB VRAM, fast training             |

docs/LLD.md

@@ -0,0 +1,96 @@
+# 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()
+```

docs/flow_diagram.drawio

@@ -0,0 +1,194 @@
+<mxfile host="65bd71144e">
+    <diagram name="Execution Flow" id="exec-flow">
+        <mxGraphModel dx="1363" dy="1186" grid="1" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" pageWidth="1600" pageHeight="900" math="0" shadow="0">
+            <root>
+                <mxCell id="0"/>
+                <mxCell id="1" parent="0"/>
+                <mxCell id="title" value="Music Generation LLM — Execution Flow" style="text;html=1;align=center;verticalAlign=middle;fontSize=18;fontStyle=1;" parent="1" vertex="1">
+                    <mxGeometry x="400" y="20" width="400" height="40" as="geometry"/>
+                </mxCell>
+                <mxCell id="s00" value="s00_main.py&#10;Entry Point&#10;CLI Args Parser" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#d5e8d4;strokeColor=#82b366;fontSize=12;" parent="1" vertex="1">
+                    <mxGeometry x="520" y="80" width="160" height="60" as="geometry"/>
+                </mxCell>
+                <mxCell id="s01" value="s01_config.py&#10;ModelConfig, TrainConfig&#10;DataConfig, GenConfig" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#fff2cc;strokeColor=#d6b656;fontSize=12;" parent="1" vertex="1">
+                    <mxGeometry x="520" y="180" width="160" height="60" as="geometry"/>
+                </mxCell>
+                <mxCell id="s02" value="s02_tokenizer.py&#10;REMI Tokenizer&#10;485 token vocab" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#dae8fc;strokeColor=#6c8ebf;fontSize=12;" parent="1" vertex="1">
+                    <mxGeometry x="520" y="280" width="160" height="60" as="geometry"/>
+                </mxCell>
+                <mxCell id="s03" value="s03_dataset.py&#10;HF Download → Parse&#10;→ Tokenize → DataLoader" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#dae8fc;strokeColor=#6c8ebf;fontSize=12;" parent="1" vertex="1">
+                    <mxGeometry x="520" y="380" width="160" height="60" as="geometry"/>
+                </mxCell>
+                <mxCell id="s04" value="s04_model.py&#10;MusicTransformer&#10;RoPE+GQA+SwiGLU+RMSNorm" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#e1d5e7;strokeColor=#9673a6;fontSize=12;" parent="1" vertex="1">
+                    <mxGeometry x="300" y="480" width="180" height="60" as="geometry"/>
+                </mxCell>
+                <mxCell id="s05" value="s05_trainer.py&#10;AMP + Grad Ckpt&#10;AdamW + Cosine LR" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#f8cecc;strokeColor=#b85450;fontSize=12;" parent="1" vertex="1">
+                    <mxGeometry x="300" y="580" width="180" height="60" as="geometry"/>
+                </mxCell>
+                <mxCell id="s06" value="s06_generator.py&#10;KV-Cache Decoding&#10;Top-p/k + RepPenalty" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#f8cecc;strokeColor=#b85450;fontSize=12;" parent="1" vertex="1">
+                    <mxGeometry x="720" y="580" width="180" height="60" as="geometry"/>
+                </mxCell>
+                <mxCell id="out_ckpt" value="checkpoints/&#10;best.pt" style="shape=cylinder3;whiteSpace=wrap;html=1;size=8;fillColor=#f5f5f5;strokeColor=#666666;fontSize=11;" parent="1" vertex="1">
+                    <mxGeometry x="340" y="680" width="100" height="60" as="geometry"/>
+                </mxCell>
+                <mxCell id="out_midi" value="output/&#10;generated_*.mid" style="shape=cylinder3;whiteSpace=wrap;html=1;size=8;fillColor=#f5f5f5;strokeColor=#666666;fontSize=11;" parent="1" vertex="1">
+                    <mxGeometry x="760" y="680" width="100" height="60" as="geometry"/>
+                </mxCell>
+                <mxCell id="s07" value="s07_utils.py&#10;• Logging&#10;• Memory Monitor&#10;• Seeding" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#fff2cc;strokeColor=#d6b656;fontSize=11;" parent="1" vertex="1">
+                    <mxGeometry x="80" y="380" width="140" height="80" as="geometry"/>
+                </mxCell>
+                <mxCell id="e1" parent="1" source="s00" target="s01" edge="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="e2" parent="1" source="s01" target="s02" edge="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="e3" parent="1" source="s02" target="s03" edge="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="e4" parent="1" source="s03" target="s04" edge="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="e5" value="train" style="edgeStyle=orthogonalEdgeStyle;" parent="1" source="s04" target="s05" edge="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="e6" value="generate" style="edgeStyle=orthogonalEdgeStyle;" parent="1" source="s04" target="s06" edge="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="e7" parent="1" source="s05" target="out_ckpt" edge="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="e8" parent="1" source="s06" target="out_midi" edge="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="e9" style="dashed=1;" parent="1" source="s07" target="s03" edge="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="e10" value="loads best.pt" style="dashed=1;" parent="1" source="out_ckpt" target="s06" edge="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+            </root>
+        </mxGraphModel>
+    </diagram>
+    <diagram name="Model Architecture" id="model-arch">
+        <mxGraphModel dx="477" dy="415" grid="1" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" pageWidth="1200" pageHeight="800" math="0" shadow="0">
+            <root>
+                <mxCell id="0"/>
+                <mxCell id="1" parent="0"/>
+                <mxCell id="t2" value="MusicTransformer Architecture" style="text;html=1;align=center;fontSize=16;fontStyle=1;" parent="1" vertex="1">
+                    <mxGeometry x="350" y="20" width="300" height="30" as="geometry"/>
+                </mxCell>
+                <mxCell id="inp" value="Token IDs&#10;[batch, seq_len]" style="shape=parallelogram;whiteSpace=wrap;html=1;fillColor=#d5e8d4;" parent="1" vertex="1">
+                    <mxGeometry x="420" y="70" width="160" height="40" as="geometry"/>
+                </mxCell>
+                <mxCell id="emb" value="Token Embedding&#10;(weight-tied with output)" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#dae8fc;" parent="1" vertex="1">
+                    <mxGeometry x="400" y="140" width="200" height="40" as="geometry"/>
+                </mxCell>
+                <mxCell id="block" value="" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#f5f5f5;strokeColor=#666;dashed=1;" parent="1" vertex="1">
+                    <mxGeometry x="300" y="210" width="400" height="300" as="geometry"/>
+                </mxCell>
+                <mxCell id="blabel" value="TransformerBlock × 6" style="text;html=1;align=center;fontSize=13;fontStyle=1;" parent="1" vertex="1">
+                    <mxGeometry x="400" y="215" width="200" height="25" as="geometry"/>
+                </mxCell>
+                <mxCell id="rms1" value="RMSNorm" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#fff2cc;" parent="1" vertex="1">
+                    <mxGeometry x="430" y="250" width="140" height="30" as="geometry"/>
+                </mxCell>
+                <mxCell id="gqa" value="Grouped Query Attention&#10;8Q / 4KV heads + RoPE&#10;Flash Attention / SDPA" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#e1d5e7;" parent="1" vertex="1">
+                    <mxGeometry x="390" y="295" width="220" height="50" as="geometry"/>
+                </mxCell>
+                <mxCell id="res1" value="+ Residual" style="ellipse;whiteSpace=wrap;html=1;fillColor=#d5e8d4;" parent="1" vertex="1">
+                    <mxGeometry x="450" y="355" width="100" height="30" as="geometry"/>
+                </mxCell>
+                <mxCell id="rms2" value="RMSNorm" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#fff2cc;" parent="1" vertex="1">
+                    <mxGeometry x="430" y="395" width="140" height="30" as="geometry"/>
+                </mxCell>
+                <mxCell id="ffn" value="SwiGLU FFN&#10;SiLU(W1·x) ⊙ W3·x → W2" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#f8cecc;" parent="1" vertex="1">
+                    <mxGeometry x="400" y="435" width="200" height="40" as="geometry"/>
+                </mxCell>
+                <mxCell id="res2" value="+ Residual" style="ellipse;whiteSpace=wrap;html=1;fillColor=#d5e8d4;" parent="1" vertex="1">
+                    <mxGeometry x="450" y="485" width="100" height="25" as="geometry"/>
+                </mxCell>
+                <mxCell id="fnorm" value="RMSNorm (final)" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#fff2cc;" parent="1" vertex="1">
+                    <mxGeometry x="430" y="540" width="140" height="30" as="geometry"/>
+                </mxCell>
+                <mxCell id="proj" value="Linear Projection → Vocab" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#dae8fc;" parent="1" vertex="1">
+                    <mxGeometry x="400" y="590" width="200" height="35" as="geometry"/>
+                </mxCell>
+                <mxCell id="out" value="Logits [batch, seq, 485]" style="shape=parallelogram;whiteSpace=wrap;html=1;fillColor=#d5e8d4;" parent="1" vertex="1">
+                    <mxGeometry x="420" y="650" width="160" height="40" as="geometry"/>
+                </mxCell>
+                <mxCell id="a1" parent="1" source="inp" target="emb" edge="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="a2" parent="1" source="emb" target="rms1" edge="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="a3" parent="1" source="rms1" target="gqa" edge="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="a4" parent="1" source="gqa" target="res1" edge="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="a5" parent="1" source="res1" target="rms2" edge="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="a6" parent="1" source="rms2" target="ffn" edge="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="a7" parent="1" source="ffn" target="res2" edge="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="a8" parent="1" source="res2" target="fnorm" edge="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="a9" parent="1" source="fnorm" target="proj" edge="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="a10" parent="1" source="proj" target="out" edge="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+            </root>
+        </mxGraphModel>
+    </diagram>
+    <diagram name="UML Class Diagram" id="uml-class">
+        <mxGraphModel dx="1400" dy="900" grid="1" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" pageWidth="1400" pageHeight="900" math="0" shadow="0">
+            <root>
+                <mxCell id="0"/>
+                <mxCell id="1" parent="0"/>
+                <!-- MusicTokenizer -->
+                <mxCell id="tok" value="MusicTokenizer&#10;─────────────&#10;- vocab_size: int&#10;- pad_id: int&#10;- bos_id: int&#10;- eos_id: int&#10;─────────────&#10;+ midi_to_tokens(midi, max_len)&#10;+ tokens_to_midi(tokens)&#10;+ decode_token(token_id)&#10;+ save(path) / load(path)" style="shape=mxgraph.uml25.class;whiteSpace=wrap;html=1;align=left;spacingLeft=8;fontSize=11;fillColor=#dae8fc;" vertex="1" parent="1">
+                    <mxGeometry x="40" y="100" width="260" height="200" as="geometry"/>
+                </mxCell>
+                <!-- MusicTransformer -->
+                <mxCell id="model" value="MusicTransformer&#10;─────────────&#10;- config: ModelConfig&#10;- token_emb: Embedding&#10;- layers: ModuleList[TransformerBlock]&#10;- norm: RMSNorm&#10;- output: Linear&#10;- freqs: Tensor (RoPE)&#10;─────────────&#10;+ forward(input_ids, targets, use_cache)&#10;+ reset_caches()&#10;+ count_parameters()&#10;+ from_config(config)" style="shape=mxgraph.uml25.class;whiteSpace=wrap;html=1;align=left;spacingLeft=8;fontSize=11;fillColor=#e1d5e7;" vertex="1" parent="1">
+                    <mxGeometry x="540" y="60" width="300" height="240" as="geometry"/>
+                </mxCell>
+                <!-- TransformerBlock -->
+                <mxCell id="tblock" value="TransformerBlock&#10;─────────────&#10;- attention: GroupedQueryAttention&#10;- feed_forward: SwiGLU&#10;- norm1, norm2: RMSNorm&#10;─────────────&#10;+ forward(x, freqs, mask, use_cache)" style="shape=mxgraph.uml25.class;whiteSpace=wrap;html=1;align=left;spacingLeft=8;fontSize=11;fillColor=#f5f5f5;" vertex="1" parent="1">
+                    <mxGeometry x="560" y="370" width="260" height="140" as="geometry"/>
+                </mxCell>
+                <!-- Trainer -->
+                <mxCell id="trainer" value="Trainer&#10;─────────────&#10;- model: MusicTransformer&#10;- optimizer: AdamW&#10;- scheduler: CosineWarmup&#10;- scaler: GradScaler&#10;─────────────&#10;+ train()&#10;+ _train_epoch(epoch)&#10;+ _validate()&#10;+ _save_checkpoint(name)&#10;+ load_checkpoint(path)" style="shape=mxgraph.uml25.class;whiteSpace=wrap;html=1;align=left;spacingLeft=8;fontSize=11;fillColor=#f8cecc;" vertex="1" parent="1">
+                    <mxGeometry x="960" y="60" width="240" height="220" as="geometry"/>
+                </mxCell>
+                <!-- Generator functions -->
+                <mxCell id="gen" value="Generator (functions)&#10;─────────────&#10;+ generate(model, tokenizer, config)&#10;+ generate_midi_file(model, tok, config, path)&#10;+ top_k_top_p_filter(logits, k, p)&#10;+ apply_repetition_penalty(logits, past, pen)" style="shape=mxgraph.uml25.class;whiteSpace=wrap;html=1;align=left;spacingLeft=8;fontSize=11;fillColor=#f8cecc;" vertex="1" parent="1">
+                    <mxGeometry x="960" y="340" width="320" height="120" as="geometry"/>
+                </mxCell>
+                <!-- Relationships -->
+                <mxCell id="r1" value="uses" style="endArrow=open;dashed=1;" edge="1" source="trainer" target="model" parent="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="r2" value="contains ×6" style="endArrow=diamond;endFill=1;" edge="1" source="tblock" target="model" parent="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="r3" value="uses" style="endArrow=open;dashed=1;" edge="1" source="gen" target="model" parent="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+                <mxCell id="r4" value="uses" style="endArrow=open;dashed=1;" edge="1" source="gen" target="tok" parent="1">
+                    <mxGeometry relative="1" as="geometry"/>
+                </mxCell>
+            </root>
+        </mxGraphModel>
+    </diagram>
+</mxfile>

inference_config.xml

@@ -0,0 +1,25 @@
+<?xml version="1.0" encoding="utf-8"?>
+<generation>
+    <!-- Sampling Hyperparameters -->
+    <settings>
+        <temperature>0.85</temperature>
+        <top_k>40</top_k>
+        <top_p>0.92</top_p>
+        <repetition_penalty>1.15</repetition_penalty>
+        <max_tokens>512</max_tokens>
+        <seed>42</seed>
+    </settings>
+
+    <!-- Seed notes to prompt the model with.
+         pitch: MIDI note (60 = Middle C)
+         velocity: note loudness (1 to 127)
+         duration_ms: note length in milliseconds
+         delay_ms: pause time before starting the note (0 for chords)
+    -->
+    <prompt>
+        <note pitch="60" velocity="100" duration_ms="300" delay_ms="0"/> <!-- C4 -->
+        <note pitch="64" velocity="100" duration_ms="300" delay_ms="0"/> <!-- E4 -->
+        <note pitch="67" velocity="100" duration_ms="300" delay_ms="0"/> <!-- G4 -->
+        <note pitch="72" velocity="100" duration_ms="600" delay_ms="300"/> <!-- C5 -->
+    </prompt>
+</generation>

requirements.txt

@@ -0,0 +1,10 @@
+torch>=2.1.0
+miditok>=3.0.0
+pretty_midi>=0.2.10
+datasets>=2.14.0
+numpy>=1.24.0
+tqdm>=4.65.0
+tensorboard>=2.14.0
+matplotlib>=3.7.0
+scipy>=1.11.0
+mido>=1.3.0

run.py

@@ -0,0 +1,68 @@
+#!/usr/bin/env python3
+"""Standalone script: download dataset, train model, generate music."""
+import subprocess
+import sys
+import os
+import glob
+
+BASE = os.path.dirname(os.path.abspath(__file__))
+DATA_DIR = os.path.join(BASE, "data")
+MIDI_DIR = os.path.join(DATA_DIR, "midi_files")
+
+def download_dataset():
+    if os.path.exists(MIDI_DIR) and glob.glob(os.path.join(MIDI_DIR, "**", "*.mid"), recursive=True):
+        print(f"Dataset already exists at {MIDI_DIR}")
+        return
+    
+    os.makedirs(DATA_DIR, exist_ok=True)
+    
+    print("Downloading MIDI dataset via git clone...")
+    env = os.environ.copy()
+    env["GIT_LFS_SKIP_SMUDGE"] = "1"
+    
+    result = subprocess.run(
+        ["git", "clone", "--depth", "1",
+         "https://huggingface.co/datasets/drengskapur/midi-classical-music",
+         MIDI_DIR],
+        env=env,
+        capture_output=True,
+        text=True,
+    )
+    
+    if result.returncode != 0:
+        print(f"Git clone failed: {result.stderr}")
+        # Fallback: use Python to download
+        print("Trying Python download fallback...")
+        from huggingface_hub import snapshot_download
+        snapshot_download(
+            repo_id="drengskapur/midi-classical-music",
+            repo_type="dataset",
+            local_dir=MIDI_DIR,
+            allow_patterns=["*.mid"],
+        )
+    
+    midi_files = glob.glob(os.path.join(MIDI_DIR, "**", "*.mid"), recursive=True)
+    print(f"Downloaded {len(midi_files)} MIDI files")
+
+def main():
+    # Step 1: Download
+    download_dataset()
+    
+    # Remove stale cache
+    cache = os.path.join(DATA_DIR, "tokenized_cache.pkl")
+    if os.path.exists(cache):
+        os.remove(cache)
+        print("Removed stale cache")
+    
+    # Step 2: Train + Generate
+    sys.path.insert(0, BASE)
+    os.chdir(BASE)
+    
+    args = sys.argv[1:] if len(sys.argv) > 1 else ["train+generate", "--epochs", "20", "--batch-size", "4", "--seq-len", "512"]
+    sys.argv = ["s00_main"] + args
+    
+    from src.s00_main import main as run_main
+    run_main()
+
+if __name__ == "__main__":
+    main()

runs/events.out.tfevents.1780546479.pop-os.1489904.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:05a052723dcdaec46d99857c6081329ce2d450a0d6e4b7b5a29c1bacd6eab1c4
+size 88

runs/events.out.tfevents.1780547909.pop-os.1514727.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a301b83b83b19017e42151abb022093b5a651ae5ad1c30ca5098e09a4a59c18d
+size 88

runs/events.out.tfevents.1780548128.pop-os.1518614.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aa823d6e7e3a241bee1d8620d39bc4ebb5e975e3c7bdecde79e9d9819c97dcb8
+size 288

runs/loss_train/events.out.tfevents.1780548163.pop-os.1518614.1

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cebdba9eeb972ca737552a0dedff012e34c98f55f7b95764354466af7575d79f
+size 298

runs/loss_val/events.out.tfevents.1780548163.pop-os.1518614.2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c851a9631104d1a512019f516518183199f4d6ceb0e9ba9dafdf0be4d363e6e8
+size 298

scripts/download_data.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+# Download and prepare the dataset
+set -e
+cd "$(dirname "$0")/.."
+echo "Installing dependencies..."
+pip install -r requirements.txt
+echo "Dataset will be downloaded automatically during training via HuggingFace datasets."
+echo "Dataset: drengskapur/midi-classical-music (~50MB, 4.8K MIDI files)"
+echo "Done!"

scripts/generate.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+# Generate music from trained model
+set -e
+cd "$(dirname "$0")/.."
+
+python -m src.s00_main generate \
+    --temperature ${TEMP:-0.85} \
+    --max-tokens ${MAX_TOKENS:-1024}

scripts/train.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+# Train the music generation model
+set -e
+cd "$(dirname "$0")/.."
+
+python -m src.s00_main train \
+    --epochs ${EPOCHS:-50} \
+    --batch-size ${BATCH_SIZE:-8} \
+    --lr ${LR:-3e-4} \
+    --seq-len ${SEQ_LEN:-1024}

src/01_config.py

@@ -0,0 +1,92 @@
+"""
+Configuration for the Music Generation LLM.
+Tuned for constrained hardware (<=8GB VRAM, <=16GB RAM).
+"""
+from dataclasses import dataclass, field
+from pathlib import Path
+import torch
+
+
+@dataclass
+class ModelConfig:
+    vocab_size: int = 0  # Set dynamically from tokenizer
+    dim: int = 256
+    n_layers: int = 6
+    n_heads: int = 8
+    n_kv_heads: int = 4  # Grouped Query Attention: fewer KV heads saves memory
+    max_seq_len: int = 1024
+    hidden_dim: int = 0  # Auto-calculated as 4 * dim * 2/3 rounded to multiple of 64
+    dropout: float = 0.1
+    rope_theta: float = 10000.0
+
+    def __post_init__(self):
+        if self.hidden_dim == 0:
+            # SwiGLU hidden dim: 4 * dim * 2/3 (LLaMA convention)
+            self.hidden_dim = int(2 * (4 * self.dim) / 3)
+            # Round to nearest multiple of 64 for hardware efficiency
+            self.hidden_dim = 64 * ((self.hidden_dim + 63) // 64)
+
+
+@dataclass
+class TrainConfig:
+    batch_size: int = 8
+    grad_accum_steps: int = 4  # Effective batch = 32
+    learning_rate: float = 3e-4
+    weight_decay: float = 0.1
+    max_epochs: int = 50
+    warmup_steps: int = 200
+    max_grad_norm: float = 1.0
+    use_amp: bool = True  # Mixed precision to save memory
+    grad_checkpoint: bool = True  # Gradient checkpointing for OOM prevention
+    eval_interval: int = 500
+    save_interval: int = 1000
+    log_interval: int = 50
+    patience: int = 10  # Early stopping patience (epochs)
+    min_delta: float = 0.001  # Minimum improvement for early stopping
+    num_workers: int = 2  # DataLoader workers (low for constrained RAM)
+    pin_memory: bool = True
+    prefetch_factor: int = 2
+
+
+@dataclass
+class DataConfig:
+    dataset_name: str = "drengskapur/midi-classical-music"
+    max_seq_len: int = 1024
+    train_split: float = 0.9
+    val_split: float = 0.1
+    tokenizer_params: str = "REMI"  # REMI tokenization — SOTA for symbolic music
+
+
+@dataclass
+class GenConfig:
+    temperature: float = 0.85
+    top_k: int = 40
+    top_p: float = 0.92
+    max_tokens: int = 1024
+    repetition_penalty: float = 1.15
+    seed: int = 42
+
+
+@dataclass
+class PathConfig:
+    base_dir: Path = field(default_factory=lambda: Path(__file__).resolve().parent.parent)
+    data_dir: Path = field(init=False)
+    checkpoint_dir: Path = field(init=False)
+    output_dir: Path = field(init=False)
+    log_dir: Path = field(init=False)
+    tokenizer_path: Path = field(init=False)
+
+    def __post_init__(self):
+        self.data_dir = self.base_dir / "data"
+        self.checkpoint_dir = self.base_dir / "checkpoints"
+        self.output_dir = self.base_dir / "output"
+        self.log_dir = self.base_dir / "runs"
+        self.tokenizer_path = self.data_dir / "tokenizer.json"
+        for d in [self.data_dir, self.checkpoint_dir, self.output_dir, self.log_dir]:
+            d.mkdir(parents=True, exist_ok=True)
+
+
+def get_device() -> torch.device:
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")

src/02_tokenizer.py

@@ -0,0 +1,250 @@
+"""
+MIDI Tokenizer using REMI (REvamped MIDI-derived) representation.
+State-of-the-art tokenization for symbolic music generation.
+Handles: Note On/Off, Velocity, Time Shift, Tempo, Time Signature.
+"""
+import json
+import logging
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+# Special token IDs
+PAD_TOKEN = 0
+BOS_TOKEN = 1
+EOS_TOKEN = 2
+SEP_TOKEN = 3
+
+# Event type offsets (after special tokens)
+SPECIAL_OFFSET = 4
+
+# REMI vocabulary layout:
+# [PAD, BOS, EOS, SEP, NoteOn_0..127, NoteOff_0..127, Velocity_0..31,
+#  TimeShift_0..99, Tempo_0..59, Position_0..31, Bar]
+NOTE_ON_OFFSET = SPECIAL_OFFSET
+NOTE_ON_COUNT = 128
+NOTE_OFF_OFFSET = NOTE_ON_OFFSET + NOTE_ON_COUNT
+NOTE_OFF_COUNT = 128
+VELOCITY_OFFSET = NOTE_OFF_OFFSET + NOTE_OFF_COUNT
+VELOCITY_COUNT = 32  # Quantized to 32 bins
+TIMESHIFT_OFFSET = VELOCITY_OFFSET + VELOCITY_COUNT
+TIMESHIFT_COUNT = 100  # 10ms to 1000ms in 10ms steps
+TEMPO_OFFSET = TIMESHIFT_OFFSET + TIMESHIFT_COUNT
+TEMPO_COUNT = 60  # 40-200 BPM quantized
+POSITION_OFFSET = TEMPO_OFFSET + TEMPO_COUNT
+POSITION_COUNT = 32  # 32 positions per bar (supports up to 32nd notes)
+BAR_OFFSET = POSITION_OFFSET + POSITION_COUNT
+BAR_COUNT = 1
+
+VOCAB_SIZE = BAR_OFFSET + BAR_COUNT
+
+
+class MusicTokenizer:
+    """Efficient REMI tokenizer for MIDI to token conversion."""
+
+    def __init__(self):
+        self.vocab_size = VOCAB_SIZE
+        self.pad_id = PAD_TOKEN
+        self.bos_id = BOS_TOKEN
+        self.eos_id = EOS_TOKEN
+
+    def note_on_token(self, pitch: int) -> int:
+        return NOTE_ON_OFFSET + max(0, min(127, pitch))
+
+    def note_off_token(self, pitch: int) -> int:
+        return NOTE_OFF_OFFSET + max(0, min(127, pitch))
+
+    def velocity_token(self, velocity: int) -> int:
+        # Quantize 0-127 to 0-31 bins
+        return VELOCITY_OFFSET + min(31, velocity // 4)
+
+    def timeshift_token(self, ms: float) -> int:
+        # Quantize to 10ms steps, capped at 1000ms
+        idx = max(0, min(99, int(ms / 10)))
+        return TIMESHIFT_OFFSET + idx
+
+    def tempo_token(self, bpm: float) -> int:
+        # Map BPM range 40-200 to 0-59
+        idx = max(0, min(59, int((bpm - 40) / (160 / 59))))
+        return TEMPO_OFFSET + idx
+
+    def position_token(self, pos: int) -> int:
+        return POSITION_OFFSET + max(0, min(31, pos))
+
+    def bar_token(self) -> int:
+        return BAR_OFFSET
+
+    def decode_token(self, token_id: int) -> dict:
+        """Decode a token ID back to its event type and value."""
+        if token_id == PAD_TOKEN:
+            return {"type": "PAD", "value": 0}
+        if token_id == BOS_TOKEN:
+            return {"type": "BOS", "value": 0}
+        if token_id == EOS_TOKEN:
+            return {"type": "EOS", "value": 0}
+        if token_id == SEP_TOKEN:
+            return {"type": "SEP", "value": 0}
+        if NOTE_ON_OFFSET <= token_id < NOTE_OFF_OFFSET:
+            return {"type": "NoteOn", "value": token_id - NOTE_ON_OFFSET}
+        if NOTE_OFF_OFFSET <= token_id < VELOCITY_OFFSET:
+            return {"type": "NoteOff", "value": token_id - NOTE_OFF_OFFSET}
+        if VELOCITY_OFFSET <= token_id < TIMESHIFT_OFFSET:
+            return {"type": "Velocity", "value": (token_id - VELOCITY_OFFSET) * 4}
+        if TIMESHIFT_OFFSET <= token_id < TEMPO_OFFSET:
+            return {"type": "TimeShift", "value": (token_id - TIMESHIFT_OFFSET) * 10}
+        if TEMPO_OFFSET <= token_id < POSITION_OFFSET:
+            return {"type": "Tempo", "value": 40 + (token_id - TEMPO_OFFSET) * (160 / 59)}
+        if POSITION_OFFSET <= token_id < BAR_OFFSET:
+            return {"type": "Position", "value": token_id - POSITION_OFFSET}
+        if token_id == BAR_OFFSET:
+            return {"type": "Bar", "value": 0}
+        return {"type": "Unknown", "value": token_id}
+
+    def midi_to_tokens(self, midi_obj, max_len: Optional[int] = None) -> list[int]:
+        """
+        Convert a pretty_midi.PrettyMIDI object to REMI token sequence.
+        Uses note-level events sorted by onset time.
+        """
+        tokens = [self.bos_id]
+
+        # Collect all notes across instruments
+        all_notes = []
+        for inst in midi_obj.instruments:
+            if inst.is_drum:
+                continue
+            for note in inst.notes:
+                all_notes.append(note)
+
+        if not all_notes:
+            tokens.append(self.eos_id)
+            return tokens
+
+        # Sort by start time, then by pitch
+        all_notes.sort(key=lambda n: (n.start, n.pitch))
+
+        # Get tempo changes
+        tempos = midi_obj.get_tempo_changes()
+        current_tempo = 120.0
+        if len(tempos[1]) > 0:
+            current_tempo = tempos[1][0]
+        tokens.append(self.tempo_token(current_tempo))
+
+        # Compute bar duration
+        bar_duration = 60.0 / current_tempo * 4  # Assume 4/4
+        current_bar = 0
+        tokens.append(self.bar_token())
+
+        prev_time = 0.0
+        for note in all_notes:
+            # Bar tracking
+            note_bar = int(note.start / bar_duration)
+            while current_bar < note_bar:
+                current_bar += 1
+                tokens.append(self.bar_token())
+
+            # Time shift from previous event
+            dt = note.start - prev_time
+            if dt > 0:
+                # Break into chunks of max 1000ms
+                while dt > 1.0:
+                    tokens.append(self.timeshift_token(1000))
+                    dt -= 1.0
+                if dt > 0.005:  # Ignore < 5ms
+                    tokens.append(self.timeshift_token(dt * 1000))
+
+            # Position within bar
+            pos_in_bar = (note.start % bar_duration) / bar_duration
+            pos_idx = int(pos_in_bar * 32)
+            tokens.append(self.position_token(pos_idx))
+
+            # Velocity then NoteOn
+            tokens.append(self.velocity_token(note.velocity))
+            tokens.append(self.note_on_token(note.pitch))
+
+            # Note duration as timeshift + NoteOff
+            dur = note.end - note.start
+            if dur > 0:
+                while dur > 1.0:
+                    tokens.append(self.timeshift_token(1000))
+                    dur -= 1.0
+                if dur > 0.005:
+                    tokens.append(self.timeshift_token(dur * 1000))
+            tokens.append(self.note_off_token(note.pitch))
+
+            prev_time = note.start
+
+            if max_len and len(tokens) >= max_len - 1:
+                break
+
+        tokens.append(self.eos_id)
+
+        if max_len:
+            tokens = tokens[:max_len]
+
+        return tokens
+
+    def tokens_to_midi(self, tokens: list[int]):
+        """Convert REMI tokens back to a PrettyMIDI object."""
+        import pretty_midi
+
+        midi = pretty_midi.PrettyMIDI(initial_tempo=120.0)
+        inst = pretty_midi.Instrument(program=0, name="Piano")
+
+        current_time = 0.0
+        current_velocity = 80
+        active_notes = {}  # pitch -> (start_time, velocity)
+
+        for token_id in tokens:
+            event = self.decode_token(token_id)
+            etype = event["type"]
+            val = event["value"]
+
+            if etype in ("PAD", "BOS", "EOS", "SEP", "Bar", "Position"):
+                continue
+            elif etype == "Tempo":
+                pass  # Could adjust timing but simpler to ignore
+            elif etype == "TimeShift":
+                current_time += val / 1000.0
+            elif etype == "Velocity":
+                current_velocity = max(1, min(127, val))
+            elif etype == "NoteOn":
+                active_notes[val] = (current_time, current_velocity)
+            elif etype == "NoteOff":
+                if val in active_notes:
+                    start, vel = active_notes.pop(val)
+                    if current_time > start:
+                        note = pretty_midi.Note(
+                            velocity=vel,
+                            pitch=val,
+                            start=start,
+                            end=current_time,
+                        )
+                        inst.notes.append(note)
+
+        # Close any remaining active notes
+        for pitch, (start, vel) in active_notes.items():
+            note = pretty_midi.Note(
+                velocity=vel, pitch=pitch, start=start, end=current_time + 0.5
+            )
+            inst.notes.append(note)
+
+        midi.instruments.append(inst)
+        return midi
+
+    def save(self, path: Path):
+        data = {"vocab_size": self.vocab_size}
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with open(path, "w") as f:
+            json.dump(data, f)
+
+    @classmethod
+    def load(cls, path: Path) -> "MusicTokenizer":
+        tok = cls()
+        if path.exists():
+            with open(path) as f:
+                data = json.load(f)
+                tok.vocab_size = data.get("vocab_size", VOCAB_SIZE)
+        return tok

src/03_dataset.py

@@ -0,0 +1,186 @@
+"""
+Data pipeline: downloads MIDI dataset, tokenizes, creates PyTorch DataLoaders.
+Uses HuggingFace datasets for efficient streaming + caching.
+Memory-efficient: processes files lazily, doesn't hold entire dataset in RAM.
+"""
+import logging
+import pickle
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset, DataLoader
+
+from src.s01_config import DataConfig, PathConfig, TrainConfig
+from src.s02_tokenizer import MusicTokenizer
+
+logger = logging.getLogger(__name__)
+
+
+class MidiTokenDataset(Dataset):
+    """
+    PyTorch dataset of pre-tokenized MIDI sequences.
+    Stores token IDs as memory-mapped numpy arrays for RAM efficiency.
+    """
+
+    def __init__(self, token_sequences: list[list[int]], max_seq_len: int, pad_id: int = 0):
+        self.max_seq_len = max_seq_len
+        self.pad_id = pad_id
+        # Filter out empty or too-short sequences
+        self.sequences = [s for s in token_sequences if len(s) >= 10]
+        logger.info(f"Dataset: {len(self.sequences)} sequences, max_len={max_seq_len}")
+
+    def __len__(self):
+        return len(self.sequences)
+
+    def __getitem__(self, idx):
+        seq = self.sequences[idx]
+
+        # For training: input = seq[:-1], target = seq[1:]
+        if len(seq) > self.max_seq_len + 1:
+            # Random crop for data augmentation
+            start = np.random.randint(0, len(seq) - self.max_seq_len)
+            seq = seq[start : start + self.max_seq_len + 1]
+
+        input_ids = seq[:-1]
+        target_ids = seq[1:]
+
+        # Pad to max_seq_len
+        pad_len = self.max_seq_len - len(input_ids)
+        if pad_len > 0:
+            input_ids = input_ids + [self.pad_id] * pad_len
+            target_ids = target_ids + [self.pad_id] * pad_len
+
+        return (
+            torch.tensor(input_ids, dtype=torch.long),
+            torch.tensor(target_ids, dtype=torch.long),
+        )
+
+
+def download_and_tokenize(
+    data_config: DataConfig,
+    path_config: PathConfig,
+    tokenizer: MusicTokenizer,
+) -> tuple[list[list[int]], list[list[int]]]:
+    """
+    Download MIDI dataset from HuggingFace and tokenize all files.
+    Returns (train_sequences, val_sequences).
+    Caches tokenized data to disk for fast reload.
+    """
+    cache_path = path_config.data_dir / "tokenized_cache.pkl"
+
+    if cache_path.exists():
+        logger.info("Loading tokenized data from cache...")
+        with open(cache_path, "rb") as f:
+            data = pickle.load(f)
+        return data["train"], data["val"]
+
+    logger.info(f"Downloading dataset: {data_config.dataset_name}")
+    from datasets import load_dataset
+    import pretty_midi
+    import io
+    import tempfile
+    import os
+
+    ds = load_dataset(data_config.dataset_name, split="train", trust_remote_code=True)
+
+    all_sequences = []
+    errors = 0
+
+    logger.info(f"Tokenizing {len(ds)} MIDI files...")
+    for i, item in enumerate(ds):
+        try:
+            midi_bytes = item.get("midi") or item.get("audio") or item.get("file")
+            if midi_bytes is None:
+                # Try getting the bytes from any binary column
+                for key, val in item.items():
+                    if isinstance(val, (bytes, dict)):
+                        if isinstance(val, dict) and "bytes" in val:
+                            midi_bytes = val["bytes"]
+                            break
+                        elif isinstance(val, bytes):
+                            midi_bytes = val
+                            break
+
+            if midi_bytes is None:
+                errors += 1
+                continue
+
+            if isinstance(midi_bytes, bytes):
+                # Write to temp file since pretty_midi needs a file path
+                with tempfile.NamedTemporaryFile(suffix=".mid", delete=False) as tmp:
+                    tmp.write(midi_bytes)
+                    tmp_path = tmp.name
+                try:
+                    midi = pretty_midi.PrettyMIDI(tmp_path)
+                    tokens = tokenizer.midi_to_tokens(midi, max_len=data_config.max_seq_len + 1)
+                    if len(tokens) >= 20:
+                        all_sequences.append(tokens)
+                finally:
+                    os.unlink(tmp_path)
+            elif isinstance(midi_bytes, str):
+                # It's a file path
+                midi = pretty_midi.PrettyMIDI(midi_bytes)
+                tokens = tokenizer.midi_to_tokens(midi, max_len=data_config.max_seq_len + 1)
+                if len(tokens) >= 20:
+                    all_sequences.append(tokens)
+
+        except Exception as e:
+            errors += 1
+            if errors <= 5:
+                logger.warning(f"Error processing item {i}: {e}")
+
+        if (i + 1) % 200 == 0:
+            logger.info(f"  Processed {i+1}/{len(ds)}, valid={len(all_sequences)}, errors={errors}")
+
+    logger.info(f"Tokenization complete: {len(all_sequences)} sequences, {errors} errors")
+
+    # Split into train/val
+    np.random.seed(42)
+    indices = np.random.permutation(len(all_sequences))
+    split = int(len(all_sequences) * data_config.train_split)
+
+    train_seqs = [all_sequences[i] for i in indices[:split]]
+    val_seqs = [all_sequences[i] for i in indices[split:]]
+
+    # Cache to disk
+    with open(cache_path, "wb") as f:
+        pickle.dump({"train": train_seqs, "val": val_seqs}, f)
+    logger.info(f"Cached tokenized data: train={len(train_seqs)}, val={len(val_seqs)}")
+
+    return train_seqs, val_seqs
+
+
+def create_dataloaders(
+    data_config: DataConfig,
+    train_config: TrainConfig,
+    path_config: PathConfig,
+    tokenizer: MusicTokenizer,
+) -> tuple[DataLoader, DataLoader]:
+    """Create train and validation DataLoaders."""
+    train_seqs, val_seqs = download_and_tokenize(data_config, path_config, tokenizer)
+
+    train_ds = MidiTokenDataset(train_seqs, data_config.max_seq_len, tokenizer.pad_id)
+    val_ds = MidiTokenDataset(val_seqs, data_config.max_seq_len, tokenizer.pad_id)
+
+    train_loader = DataLoader(
+        train_ds,
+        batch_size=train_config.batch_size,
+        shuffle=True,
+        num_workers=train_config.num_workers,
+        pin_memory=train_config.pin_memory,
+        prefetch_factor=train_config.prefetch_factor,
+        drop_last=True,
+    )
+    val_loader = DataLoader(
+        val_ds,
+        batch_size=train_config.batch_size,
+        shuffle=False,
+        num_workers=train_config.num_workers,
+        pin_memory=train_config.pin_memory,
+        prefetch_factor=train_config.prefetch_factor,
+        drop_last=False,
+    )
+
+    return train_loader, val_loader

src/s00_main.py

@@ -0,0 +1,167 @@
+"""
+Main entry point for Music Generation LLM.
+Orchestrates: config → data → model → train → generate.
+
+Usage:
+    python -m src.s00_main train          # Train the model
+    python -m src.s00_main generate       # Generate music from trained model
+    python -m src.s00_main train+generate # Train then generate
+"""
+import argparse
+import logging
+import sys
+from pathlib import Path
+
+import torch
+
+# Add project root to path
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from src.s01_config import ModelConfig, TrainConfig, DataConfig, GenConfig, PathConfig, get_device
+from src.s02_tokenizer import MusicTokenizer
+from src.s03_dataset import create_dataloaders
+from src.s04_model import MusicTransformer
+from src.s05_trainer import Trainer
+from src.s06_generator import generate_midi_file
+from src.s07_utils import setup_logging, set_seed, log_memory_usage, clear_memory
+
+logger = logging.getLogger(__name__)
+
+
+def train_pipeline(
+    model_config: ModelConfig,
+    train_config: TrainConfig,
+    data_config: DataConfig,
+    path_config: PathConfig,
+    tokenizer: MusicTokenizer,
+):
+    """Full training pipeline."""
+    logger.info("=" * 60)
+    logger.info("MUSIC GENERATION LLM — TRAINING")
+    logger.info("=" * 60)
+
+    # Create data loaders
+    logger.info("Preparing data...")
+    train_loader, val_loader = create_dataloaders(
+        data_config, train_config, path_config, tokenizer
+    )
+
+    # Build model
+    model_config.vocab_size = tokenizer.vocab_size
+    model = MusicTransformer.from_config(model_config)
+    logger.info(f"Model: {model.count_parameters():,} parameters")
+    log_memory_usage("Pre-training")
+
+    # Train
+    trainer = Trainer(model, train_loader, val_loader, train_config, path_config)
+    trainer.train()
+
+    # Save tokenizer
+    tokenizer.save(path_config.tokenizer_path)
+    log_memory_usage("Post-training")
+    clear_memory()
+    return model
+
+
+def generate_pipeline(
+    model_config: ModelConfig,
+    gen_config: GenConfig,
+    path_config: PathConfig,
+    tokenizer: MusicTokenizer,
+    model: MusicTransformer | None = None,
+):
+    """Generate music from trained model."""
+    logger.info("=" * 60)
+    logger.info("MUSIC GENERATION LLM — GENERATING")
+    logger.info("=" * 60)
+
+    if model is None:
+        best_ckpt = path_config.checkpoint_dir / "best.pt"
+        if not best_ckpt.exists():
+            logger.error(f"No checkpoint found at {best_ckpt}. Train first!")
+            return
+
+        model_config.vocab_size = tokenizer.vocab_size
+        model = MusicTransformer.from_config(model_config)
+        ckpt = torch.load(best_ckpt, map_location=get_device(), weights_only=False)
+        model.load_state_dict(ckpt["model_state_dict"])
+        logger.info(f"Loaded model from {best_ckpt}")
+
+    # Generate multiple samples
+    for i in range(3):
+        output_path = path_config.output_dir / f"generated_{i+1}.mid"
+        gen_config_i = GenConfig(
+            temperature=gen_config.temperature,
+            top_k=gen_config.top_k,
+            top_p=gen_config.top_p,
+            max_tokens=gen_config.max_tokens,
+            repetition_penalty=gen_config.repetition_penalty,
+            seed=gen_config.seed + i,
+        )
+        generate_midi_file(model, tokenizer, gen_config_i, output_path)
+        logger.info(f"Generated sample {i+1}: {output_path}")
+
+    clear_memory()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Music Generation LLM")
+    parser.add_argument(
+        "mode",
+        choices=["train", "generate", "train+generate"],
+        default="train+generate",
+        nargs="?",
+        help="Operation mode",
+    )
+    parser.add_argument("--epochs", type=int, default=None, help="Override max epochs")
+    parser.add_argument("--batch-size", type=int, default=None, help="Override batch size")
+    parser.add_argument("--lr", type=float, default=None, help="Override learning rate")
+    parser.add_argument("--seq-len", type=int, default=None, help="Override max sequence length")
+    parser.add_argument("--temperature", type=float, default=None, help="Generation temperature")
+    parser.add_argument("--max-tokens", type=int, default=None, help="Max generation tokens")
+    args = parser.parse_args()
+
+    setup_logging()
+    set_seed(42)
+
+    # Initialize configs
+    model_config = ModelConfig()
+    train_config = TrainConfig()
+    data_config = DataConfig()
+    gen_config = GenConfig()
+    path_config = PathConfig()
+
+    # Apply overrides
+    if args.epochs:
+        train_config.max_epochs = args.epochs
+    if args.batch_size:
+        train_config.batch_size = args.batch_size
+    if args.lr:
+        train_config.learning_rate = args.lr
+    if args.seq_len:
+        data_config.max_seq_len = args.seq_len
+        model_config.max_seq_len = args.seq_len
+    if args.temperature:
+        gen_config.temperature = args.temperature
+    if args.max_tokens:
+        gen_config.max_tokens = args.max_tokens
+
+    # Tokenizer
+    tokenizer = MusicTokenizer()
+    model_config.vocab_size = tokenizer.vocab_size
+
+    logger.info(f"Device: {get_device()}")
+    logger.info(f"Vocab size: {tokenizer.vocab_size}")
+    logger.info(f"Model dim: {model_config.dim}, layers: {model_config.n_layers}, "
+                f"heads: {model_config.n_heads}, kv_heads: {model_config.n_kv_heads}")
+
+    model = None
+    if "train" in args.mode:
+        model = train_pipeline(model_config, train_config, data_config, path_config, tokenizer)
+
+    if "generate" in args.mode:
+        generate_pipeline(model_config, gen_config, path_config, tokenizer, model)
+
+
+if __name__ == "__main__":
+    main()

src/s01_config.py

@@ -0,0 +1,92 @@
+"""
+Configuration for the Music Generation LLM.
+Tuned for constrained hardware (<=8GB VRAM, <=16GB RAM).
+"""
+from dataclasses import dataclass, field
+from pathlib import Path
+import torch
+
+
+@dataclass
+class ModelConfig:
+    vocab_size: int = 0  # Set dynamically from tokenizer
+    dim: int = 256
+    n_layers: int = 6
+    n_heads: int = 8
+    n_kv_heads: int = 4  # Grouped Query Attention: fewer KV heads saves memory
+    max_seq_len: int = 1024
+    hidden_dim: int = 0  # Auto-calculated as 4 * dim * 2/3 rounded to multiple of 64
+    dropout: float = 0.1
+    rope_theta: float = 10000.0
+
+    def __post_init__(self):
+        if self.hidden_dim == 0:
+            # SwiGLU hidden dim: 4 * dim * 2/3 (LLaMA convention)
+            self.hidden_dim = int(2 * (4 * self.dim) / 3)
+            # Round to nearest multiple of 64 for hardware efficiency
+            self.hidden_dim = 64 * ((self.hidden_dim + 63) // 64)
+
+
+@dataclass
+class TrainConfig:
+    batch_size: int = 8
+    grad_accum_steps: int = 4  # Effective batch = 32
+    learning_rate: float = 3e-4
+    weight_decay: float = 0.1
+    max_epochs: int = 50
+    warmup_steps: int = 200
+    max_grad_norm: float = 1.0
+    use_amp: bool = True  # Mixed precision to save memory
+    grad_checkpoint: bool = True  # Gradient checkpointing for OOM prevention
+    eval_interval: int = 500
+    save_interval: int = 1000
+    log_interval: int = 50
+    patience: int = 10  # Early stopping patience (epochs)
+    min_delta: float = 0.001  # Minimum improvement for early stopping
+    num_workers: int = 2  # DataLoader workers (low for constrained RAM)
+    pin_memory: bool = True
+    prefetch_factor: int = 2
+
+
+@dataclass
+class DataConfig:
+    dataset_name: str = "drengskapur/midi-classical-music"
+    max_seq_len: int = 1024
+    train_split: float = 0.9
+    val_split: float = 0.1
+    tokenizer_params: str = "REMI"  # REMI tokenization — SOTA for symbolic music
+
+
+@dataclass
+class GenConfig:
+    temperature: float = 0.85
+    top_k: int = 40
+    top_p: float = 0.92
+    max_tokens: int = 1024
+    repetition_penalty: float = 1.15
+    seed: int = 42
+
+
+@dataclass
+class PathConfig:
+    base_dir: Path = field(default_factory=lambda: Path(__file__).resolve().parent.parent)
+    data_dir: Path = field(init=False)
+    checkpoint_dir: Path = field(init=False)
+    output_dir: Path = field(init=False)
+    log_dir: Path = field(init=False)
+    tokenizer_path: Path = field(init=False)
+
+    def __post_init__(self):
+        self.data_dir = self.base_dir / "data"
+        self.checkpoint_dir = self.base_dir / "checkpoints"
+        self.output_dir = self.base_dir / "output"
+        self.log_dir = self.base_dir / "runs"
+        self.tokenizer_path = self.data_dir / "tokenizer.json"
+        for d in [self.data_dir, self.checkpoint_dir, self.output_dir, self.log_dir]:
+            d.mkdir(parents=True, exist_ok=True)
+
+
+def get_device() -> torch.device:
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")

src/s02_tokenizer.py

@@ -0,0 +1,250 @@
+"""
+MIDI Tokenizer using REMI (REvamped MIDI-derived) representation.
+State-of-the-art tokenization for symbolic music generation.
+Handles: Note On/Off, Velocity, Time Shift, Tempo, Time Signature.
+"""
+import json
+import logging
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+# Special token IDs
+PAD_TOKEN = 0
+BOS_TOKEN = 1
+EOS_TOKEN = 2
+SEP_TOKEN = 3
+
+# Event type offsets (after special tokens)
+SPECIAL_OFFSET = 4
+
+# REMI vocabulary layout:
+# [PAD, BOS, EOS, SEP, NoteOn_0..127, NoteOff_0..127, Velocity_0..31,
+#  TimeShift_0..99, Tempo_0..59, Position_0..31, Bar]
+NOTE_ON_OFFSET = SPECIAL_OFFSET
+NOTE_ON_COUNT = 128
+NOTE_OFF_OFFSET = NOTE_ON_OFFSET + NOTE_ON_COUNT
+NOTE_OFF_COUNT = 128
+VELOCITY_OFFSET = NOTE_OFF_OFFSET + NOTE_OFF_COUNT
+VELOCITY_COUNT = 32  # Quantized to 32 bins
+TIMESHIFT_OFFSET = VELOCITY_OFFSET + VELOCITY_COUNT
+TIMESHIFT_COUNT = 100  # 10ms to 1000ms in 10ms steps
+TEMPO_OFFSET = TIMESHIFT_OFFSET + TIMESHIFT_COUNT
+TEMPO_COUNT = 60  # 40-200 BPM quantized
+POSITION_OFFSET = TEMPO_OFFSET + TEMPO_COUNT
+POSITION_COUNT = 32  # 32 positions per bar (supports up to 32nd notes)
+BAR_OFFSET = POSITION_OFFSET + POSITION_COUNT
+BAR_COUNT = 1
+
+VOCAB_SIZE = BAR_OFFSET + BAR_COUNT
+
+
+class MusicTokenizer:
+    """Efficient REMI tokenizer for MIDI to token conversion."""
+
+    def __init__(self):
+        self.vocab_size = VOCAB_SIZE
+        self.pad_id = PAD_TOKEN
+        self.bos_id = BOS_TOKEN
+        self.eos_id = EOS_TOKEN
+
+    def note_on_token(self, pitch: int) -> int:
+        return NOTE_ON_OFFSET + max(0, min(127, pitch))
+
+    def note_off_token(self, pitch: int) -> int:
+        return NOTE_OFF_OFFSET + max(0, min(127, pitch))
+
+    def velocity_token(self, velocity: int) -> int:
+        # Quantize 0-127 to 0-31 bins
+        return VELOCITY_OFFSET + min(31, velocity // 4)
+
+    def timeshift_token(self, ms: float) -> int:
+        # Quantize to 10ms steps, capped at 1000ms
+        idx = max(0, min(99, int(ms / 10)))
+        return TIMESHIFT_OFFSET + idx
+
+    def tempo_token(self, bpm: float) -> int:
+        # Map BPM range 40-200 to 0-59
+        idx = max(0, min(59, int((bpm - 40) / (160 / 59))))
+        return TEMPO_OFFSET + idx
+
+    def position_token(self, pos: int) -> int:
+        return POSITION_OFFSET + max(0, min(31, pos))
+
+    def bar_token(self) -> int:
+        return BAR_OFFSET
+
+    def decode_token(self, token_id: int) -> dict:
+        """Decode a token ID back to its event type and value."""
+        if token_id == PAD_TOKEN:
+            return {"type": "PAD", "value": 0}
+        if token_id == BOS_TOKEN:
+            return {"type": "BOS", "value": 0}
+        if token_id == EOS_TOKEN:
+            return {"type": "EOS", "value": 0}
+        if token_id == SEP_TOKEN:
+            return {"type": "SEP", "value": 0}
+        if NOTE_ON_OFFSET <= token_id < NOTE_OFF_OFFSET:
+            return {"type": "NoteOn", "value": token_id - NOTE_ON_OFFSET}
+        if NOTE_OFF_OFFSET <= token_id < VELOCITY_OFFSET:
+            return {"type": "NoteOff", "value": token_id - NOTE_OFF_OFFSET}
+        if VELOCITY_OFFSET <= token_id < TIMESHIFT_OFFSET:
+            return {"type": "Velocity", "value": (token_id - VELOCITY_OFFSET) * 4}
+        if TIMESHIFT_OFFSET <= token_id < TEMPO_OFFSET:
+            return {"type": "TimeShift", "value": (token_id - TIMESHIFT_OFFSET) * 10}
+        if TEMPO_OFFSET <= token_id < POSITION_OFFSET:
+            return {"type": "Tempo", "value": 40 + (token_id - TEMPO_OFFSET) * (160 / 59)}
+        if POSITION_OFFSET <= token_id < BAR_OFFSET:
+            return {"type": "Position", "value": token_id - POSITION_OFFSET}
+        if token_id == BAR_OFFSET:
+            return {"type": "Bar", "value": 0}
+        return {"type": "Unknown", "value": token_id}
+
+    def midi_to_tokens(self, midi_obj, max_len: Optional[int] = None) -> list[int]:
+        """
+        Convert a pretty_midi.PrettyMIDI object to REMI token sequence.
+        Uses note-level events sorted by onset time.
+        """
+        tokens = [self.bos_id]
+
+        # Collect all notes across instruments
+        all_notes = []
+        for inst in midi_obj.instruments:
+            if inst.is_drum:
+                continue
+            for note in inst.notes:
+                all_notes.append(note)
+
+        if not all_notes:
+            tokens.append(self.eos_id)
+            return tokens
+
+        # Sort by start time, then by pitch
+        all_notes.sort(key=lambda n: (n.start, n.pitch))
+
+        # Get tempo changes
+        tempos = midi_obj.get_tempo_changes()
+        current_tempo = 120.0
+        if len(tempos[1]) > 0:
+            current_tempo = tempos[1][0]
+        tokens.append(self.tempo_token(current_tempo))
+
+        # Compute bar duration
+        bar_duration = 60.0 / current_tempo * 4  # Assume 4/4
+        current_bar = 0
+        tokens.append(self.bar_token())
+
+        prev_time = 0.0
+        for note in all_notes:
+            # Bar tracking
+            note_bar = int(note.start / bar_duration)
+            while current_bar < note_bar:
+                current_bar += 1
+                tokens.append(self.bar_token())
+
+            # Time shift from previous event
+            dt = note.start - prev_time
+            if dt > 0:
+                # Break into chunks of max 1000ms
+                while dt > 1.0:
+                    tokens.append(self.timeshift_token(1000))
+                    dt -= 1.0
+                if dt > 0.005:  # Ignore < 5ms
+                    tokens.append(self.timeshift_token(dt * 1000))
+
+            # Position within bar
+            pos_in_bar = (note.start % bar_duration) / bar_duration
+            pos_idx = int(pos_in_bar * 32)
+            tokens.append(self.position_token(pos_idx))
+
+            # Velocity then NoteOn
+            tokens.append(self.velocity_token(note.velocity))
+            tokens.append(self.note_on_token(note.pitch))
+
+            # Note duration as timeshift + NoteOff
+            dur = note.end - note.start
+            if dur > 0:
+                while dur > 1.0:
+                    tokens.append(self.timeshift_token(1000))
+                    dur -= 1.0
+                if dur > 0.005:
+                    tokens.append(self.timeshift_token(dur * 1000))
+            tokens.append(self.note_off_token(note.pitch))
+
+            prev_time = note.start
+
+            if max_len and len(tokens) >= max_len - 1:
+                break
+
+        tokens.append(self.eos_id)
+
+        if max_len:
+            tokens = tokens[:max_len]
+
+        return tokens
+
+    def tokens_to_midi(self, tokens: list[int]):
+        """Convert REMI tokens back to a PrettyMIDI object."""
+        import pretty_midi
+
+        midi = pretty_midi.PrettyMIDI(initial_tempo=120.0)
+        inst = pretty_midi.Instrument(program=0, name="Piano")
+
+        current_time = 0.0
+        current_velocity = 80
+        active_notes = {}  # pitch -> (start_time, velocity)
+
+        for token_id in tokens:
+            event = self.decode_token(token_id)
+            etype = event["type"]
+            val = event["value"]
+
+            if etype in ("PAD", "BOS", "EOS", "SEP", "Bar", "Position"):
+                continue
+            elif etype == "Tempo":
+                pass  # Could adjust timing but simpler to ignore
+            elif etype == "TimeShift":
+                current_time += val / 1000.0
+            elif etype == "Velocity":
+                current_velocity = max(1, min(127, val))
+            elif etype == "NoteOn":
+                active_notes[val] = (current_time, current_velocity)
+            elif etype == "NoteOff":
+                if val in active_notes:
+                    start, vel = active_notes.pop(val)
+                    if current_time > start:
+                        note = pretty_midi.Note(
+                            velocity=vel,
+                            pitch=val,
+                            start=start,
+                            end=current_time,
+                        )
+                        inst.notes.append(note)
+
+        # Close any remaining active notes
+        for pitch, (start, vel) in active_notes.items():
+            note = pretty_midi.Note(
+                velocity=vel, pitch=pitch, start=start, end=current_time + 0.5
+            )
+            inst.notes.append(note)
+
+        midi.instruments.append(inst)
+        return midi
+
+    def save(self, path: Path):
+        data = {"vocab_size": self.vocab_size}
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with open(path, "w") as f:
+            json.dump(data, f)
+
+    @classmethod
+    def load(cls, path: Path) -> "MusicTokenizer":
+        tok = cls()
+        if path.exists():
+            with open(path) as f:
+                data = json.load(f)
+                tok.vocab_size = data.get("vocab_size", VOCAB_SIZE)
+        return tok

src/s03_dataset.py

@@ -0,0 +1,198 @@
+"""
+Data pipeline: downloads MIDI dataset, tokenizes, creates PyTorch DataLoaders.
+Uses HuggingFace datasets for efficient streaming + caching.
+Memory-efficient: processes files lazily, doesn't hold entire dataset in RAM.
+"""
+import logging
+import os
+import pickle
+import signal
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset, DataLoader
+
+
+class _TimeoutError(Exception):
+    pass
+
+
+def _timeout_handler(signum, frame):
+    raise _TimeoutError("File processing timed out")
+
+from src.s01_config import DataConfig, PathConfig, TrainConfig
+from src.s02_tokenizer import MusicTokenizer
+
+logger = logging.getLogger(__name__)
+
+
+class MidiTokenDataset(Dataset):
+    """
+    PyTorch dataset of pre-tokenized MIDI sequences.
+    Stores token IDs as memory-mapped numpy arrays for RAM efficiency.
+    """
+
+    def __init__(self, token_sequences: list[list[int]], max_seq_len: int, pad_id: int = 0):
+        self.max_seq_len = max_seq_len
+        self.pad_id = pad_id
+        # Filter out empty or too-short sequences
+        self.sequences = [s for s in token_sequences if len(s) >= 10]
+        logger.info(f"Dataset: {len(self.sequences)} sequences, max_len={max_seq_len}")
+
+    def __len__(self):
+        return len(self.sequences)
+
+    def __getitem__(self, idx):
+        seq = self.sequences[idx]
+
+        # For training: input = seq[:-1], target = seq[1:]
+        if len(seq) > self.max_seq_len + 1:
+            # Random crop for data augmentation
+            start = np.random.randint(0, len(seq) - self.max_seq_len)
+            seq = seq[start : start + self.max_seq_len + 1]
+
+        input_ids = seq[:-1]
+        target_ids = seq[1:]
+
+        # Pad to max_seq_len
+        pad_len = self.max_seq_len - len(input_ids)
+        if pad_len > 0:
+            input_ids = input_ids + [self.pad_id] * pad_len
+            target_ids = target_ids + [self.pad_id] * pad_len
+
+        return (
+            torch.tensor(input_ids, dtype=torch.long),
+            torch.tensor(target_ids, dtype=torch.long),
+        )
+
+
+def download_and_tokenize(
+    data_config: DataConfig,
+    path_config: PathConfig,
+    tokenizer: MusicTokenizer,
+) -> tuple[list[list[int]], list[list[int]]]:
+    """
+    Download MIDI dataset from HuggingFace and tokenize all files.
+    Returns (train_sequences, val_sequences).
+    Caches tokenized data to disk for fast reload.
+    """
+    cache_path = path_config.data_dir / "tokenized_cache.pkl"
+
+    if cache_path.exists():
+        logger.info("Loading tokenized data from cache...")
+        with open(cache_path, "rb") as f:
+            data = pickle.load(f)
+        return data["train"], data["val"]
+
+    logger.info(f"Downloading dataset: {data_config.dataset_name}")
+    import pretty_midi
+    import subprocess
+    import glob
+
+    # Fast git clone instead of slow per-file snapshot_download
+    midi_dir = path_config.data_dir / "midi_files"
+    if not midi_dir.exists():
+        logger.info("Cloning MIDI dataset repo (faster than per-file download)...")
+        subprocess.run(
+            ["git", "clone", "--depth", "1",
+             f"https://huggingface.co/datasets/{data_config.dataset_name}",
+             str(midi_dir)],
+            check=True,
+        )
+    else:
+        logger.info(f"Using cached MIDI files from {midi_dir}")
+
+    # Find all MIDI files
+    midi_files = sorted(
+        glob.glob(f"{midi_dir}/**/*.mid", recursive=True)
+        + glob.glob(f"{midi_dir}/**/*.midi", recursive=True)
+        + glob.glob(f"{midi_dir}/**/*.MID", recursive=True)
+    )
+    logger.info(f"Found {len(midi_files)} MIDI files")
+
+    all_sequences = []
+    errors = 0
+
+    for i, midi_path in enumerate(midi_files):
+        try:
+            # Skip files > 100KB (large orchestral pieces cause slow processing)
+            if os.path.getsize(midi_path) > 100_000:
+                errors += 1
+                continue
+            # Skip git LFS pointer files
+            with open(midi_path, "rb") as f:
+                header = f.read(20)
+            if header.startswith(b"version https://git"):
+                errors += 1
+                continue
+            # 30s timeout per file to avoid hangs on complex MIDI
+            old_handler = signal.signal(signal.SIGALRM, _timeout_handler)
+            signal.alarm(30)
+            try:
+                midi = pretty_midi.PrettyMIDI(midi_path)
+                tokens = tokenizer.midi_to_tokens(midi, max_len=data_config.max_seq_len + 1)
+                if len(tokens) >= 20:
+                    all_sequences.append(tokens)
+            finally:
+                signal.alarm(0)
+                signal.signal(signal.SIGALRM, old_handler)
+        except (_TimeoutError, Exception) as e:
+            errors += 1
+            if errors <= 10:
+                logger.warning(f"Error processing {midi_path}: {e}")
+
+        if (i + 1) % 200 == 0:
+            logger.info(f"  Processed {i+1}/{len(midi_files)}, valid={len(all_sequences)}, errors={errors}")
+
+    logger.info(f"Tokenization complete: {len(all_sequences)} sequences, {errors} errors")
+
+    # Split into train/val
+    np.random.seed(42)
+    indices = np.random.permutation(len(all_sequences))
+    split = int(len(all_sequences) * data_config.train_split)
+
+    train_seqs = [all_sequences[i] for i in indices[:split]]
+    val_seqs = [all_sequences[i] for i in indices[split:]]
+
+    # Cache to disk
+    with open(cache_path, "wb") as f:
+        pickle.dump({"train": train_seqs, "val": val_seqs}, f)
+    logger.info(f"Cached tokenized data: train={len(train_seqs)}, val={len(val_seqs)}")
+
+    return train_seqs, val_seqs
+
+
+def create_dataloaders(
+    data_config: DataConfig,
+    train_config: TrainConfig,
+    path_config: PathConfig,
+    tokenizer: MusicTokenizer,
+) -> tuple[DataLoader, DataLoader]:
+    """Create train and validation DataLoaders."""
+    train_seqs, val_seqs = download_and_tokenize(data_config, path_config, tokenizer)
+
+    train_ds = MidiTokenDataset(train_seqs, data_config.max_seq_len, tokenizer.pad_id)
+    val_ds = MidiTokenDataset(val_seqs, data_config.max_seq_len, tokenizer.pad_id)
+
+    train_loader = DataLoader(
+        train_ds,
+        batch_size=train_config.batch_size,
+        shuffle=True,
+        num_workers=train_config.num_workers,
+        pin_memory=train_config.pin_memory,
+        prefetch_factor=train_config.prefetch_factor,
+        drop_last=True,
+    )
+    val_loader = DataLoader(
+        val_ds,
+        batch_size=train_config.batch_size,
+        shuffle=False,
+        num_workers=train_config.num_workers,
+        pin_memory=train_config.pin_memory,
+        prefetch_factor=train_config.prefetch_factor,
+        drop_last=False,
+    )
+
+    return train_loader, val_loader

src/s04_model.py

@@ -0,0 +1,310 @@
+"""
+Music Transformer Model — LLaMA-style architecture for symbolic music generation.
+
+Key innovations combined:
+- Rotary Position Embeddings (RoPE) — better long-range modeling than sinusoidal
+- RMSNorm — faster than LayerNorm, used in LLaMA/Mistral
+- SwiGLU activation — better than GELU/ReLU, used in LLaMA
+- Grouped Query Attention (GQA) — reduces KV-cache memory by sharing KV heads
+- Gradient checkpointing — cuts memory usage ~50% with ~20% speed cost
+- KV-cache — O(1) per-token inference instead of O(n)
+"""
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class RMSNorm(nn.Module):
+    """Root Mean Square Layer Normalization (faster than LayerNorm)."""
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
+        return (x.float() * norm).type_as(x) * self.weight
+
+
+def precompute_rope_freqs(dim: int, max_seq_len: int, theta: float = 10000.0) -> torch.Tensor:
+    """Precompute RoPE frequency tensor for complex exponentials."""
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
+    t = torch.arange(max_seq_len, dtype=torch.float32)
+    freqs = torch.outer(t, freqs)
+    return torch.polar(torch.ones_like(freqs), freqs)  # complex64
+
+
+def apply_rope(xq: torch.Tensor, xk: torch.Tensor, freqs: torch.Tensor):
+    """Apply rotary embeddings to query and key tensors."""
+    # Reshape to complex
+    xq_c = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
+    xk_c = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+
+    # Reshape freqs for broadcasting: (seq_len,) -> (1, seq_len, 1, head_dim//2)
+    freqs = freqs.unsqueeze(0).unsqueeze(2)
+
+    xq_out = torch.view_as_real(xq_c * freqs).flatten(-2)
+    xk_out = torch.view_as_real(xk_c * freqs).flatten(-2)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """Repeat KV heads to match query head count for GQA."""
+    if n_rep == 1:
+        return x
+    bs, seq_len, n_kv_heads, head_dim = x.shape
+    return (
+        x[:, :, :, None, :]
+        .expand(bs, seq_len, n_kv_heads, n_rep, head_dim)
+        .reshape(bs, seq_len, n_kv_heads * n_rep, head_dim)
+    )
+
+
+class GroupedQueryAttention(nn.Module):
+    """
+    Multi-head attention with Grouped Query Attention (GQA).
+    Uses fewer KV heads than Q heads to reduce memory.
+    """
+
+    def __init__(self, dim: int, n_heads: int, n_kv_heads: int, dropout: float = 0.1):
+        super().__init__()
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.n_rep = n_heads // n_kv_heads
+        self.head_dim = dim // n_heads
+
+        self.wq = nn.Linear(dim, n_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(dim, n_kv_heads * self.head_dim, bias=False)
+        self.wv = nn.Linear(dim, n_kv_heads * self.head_dim, bias=False)
+        self.wo = nn.Linear(n_heads * self.head_dim, dim, bias=False)
+        self.attn_dropout = nn.Dropout(dropout)
+        self.resid_dropout = nn.Dropout(dropout)
+
+        # KV-cache for inference
+        self.cache_k: Optional[torch.Tensor] = None
+        self.cache_v: Optional[torch.Tensor] = None
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        freqs: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        use_cache: bool = False,
+    ) -> torch.Tensor:
+        bs, seq_len, _ = x.shape
+
+        q = self.wq(x).view(bs, seq_len, self.n_heads, self.head_dim)
+        k = self.wk(x).view(bs, seq_len, self.n_kv_heads, self.head_dim)
+        v = self.wv(x).view(bs, seq_len, self.n_kv_heads, self.head_dim)
+
+        # Apply RoPE to Q and K
+        q_rope = q.view(bs, seq_len, self.n_heads, self.head_dim)
+        k_rope = k.view(bs, seq_len, self.n_kv_heads, self.head_dim)
+
+        # RoPE needs (bs, seq_len, heads, head_dim) but freqs is (seq_len, head_dim//2)
+        # Apply per-head
+        q_for_rope = q_rope.reshape(bs * self.n_heads, seq_len, self.head_dim)
+        k_for_rope = k_rope.reshape(bs * self.n_kv_heads, seq_len, self.head_dim)
+
+        # Simpler RoPE application
+        q = q.transpose(1, 2)  # (bs, n_heads, seq_len, head_dim)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        # Apply RoPE via cos/sin (more compatible than complex)
+        q, k = self._apply_rope_real(q, k, freqs)
+
+        # KV-cache for generation
+        if use_cache:
+            if self.cache_k is not None:
+                k = torch.cat([self.cache_k, k], dim=2)
+                v = torch.cat([self.cache_v, v], dim=2)
+            self.cache_k = k.detach()
+            self.cache_v = v.detach()
+
+        # GQA: repeat KV heads
+        k = repeat_kv(k.transpose(1, 2), self.n_rep).transpose(1, 2)
+        v = repeat_kv(v.transpose(1, 2), self.n_rep).transpose(1, 2)
+
+        # Scaled dot-product attention (uses Flash Attention when available)
+        scale = 1.0 / math.sqrt(self.head_dim)
+        try:
+            # PyTorch 2.0+ SDPA with memory-efficient backend
+            out = F.scaled_dot_product_attention(
+                q, k, v,
+                attn_mask=mask,
+                dropout_p=self.attn_dropout.p if self.training else 0.0,
+                is_causal=(mask is None and not use_cache),
+            )
+        except RuntimeError:
+            # Fallback for older PyTorch
+            scores = torch.matmul(q, k.transpose(-2, -1)) * scale
+            if mask is not None:
+                scores = scores + mask
+            elif not use_cache:
+                causal = torch.triu(
+                    torch.full((seq_len, seq_len), float("-inf"), device=x.device), diagonal=1
+                )
+                scores = scores + causal
+            scores = F.softmax(scores, dim=-1)
+            scores = self.attn_dropout(scores)
+            out = torch.matmul(scores, v)
+
+        out = out.transpose(1, 2).contiguous().view(bs, seq_len, -1)
+        return self.resid_dropout(self.wo(out))
+
+    def _apply_rope_real(self, q, k, freqs):
+        """Apply RoPE using real-valued sin/cos (more device-compatible)."""
+        # freqs shape: (seq_len, head_dim//2)
+        seq_len = q.shape[2]
+        freqs = freqs[:seq_len]
+
+        cos_f = freqs.cos().unsqueeze(0).unsqueeze(0)  # (1, 1, seq_len, head_dim//2)
+        sin_f = freqs.sin().unsqueeze(0).unsqueeze(0)
+
+        def rotate_half(x):
+            x1, x2 = x.chunk(2, dim=-1)
+            return torch.cat((-x2, x1), dim=-1)
+
+        q = q * cos_f.repeat(1, 1, 1, 2) + rotate_half(q) * sin_f.repeat(1, 1, 1, 2)
+        k = k * cos_f.repeat(1, 1, 1, 2) + rotate_half(k) * sin_f.repeat(1, 1, 1, 2)
+        return q, k
+
+    def reset_cache(self):
+        self.cache_k = None
+        self.cache_v = None
+
+
+class SwiGLU(nn.Module):
+    """SwiGLU activation — superior to GELU/ReLU, used in LLaMA."""
+
+    def __init__(self, dim: int, hidden_dim: int, dropout: float = 0.1):
+        super().__init__()
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
+
+
+class TransformerBlock(nn.Module):
+    """Single transformer block with pre-norm architecture."""
+
+    def __init__(self, dim: int, n_heads: int, n_kv_heads: int, hidden_dim: int, dropout: float):
+        super().__init__()
+        self.attention = GroupedQueryAttention(dim, n_heads, n_kv_heads, dropout)
+        self.feed_forward = SwiGLU(dim, hidden_dim, dropout)
+        self.norm1 = RMSNorm(dim)
+        self.norm2 = RMSNorm(dim)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        freqs: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        use_cache: bool = False,
+    ) -> torch.Tensor:
+        # Pre-norm residual connections
+        x = x + self.attention(self.norm1(x), freqs, mask, use_cache)
+        x = x + self.feed_forward(self.norm2(x))
+        return x
+
+
+class MusicTransformer(nn.Module):
+    """
+    LLaMA-style Transformer for music generation.
+    Combines: RoPE + GQA + SwiGLU + RMSNorm + gradient checkpointing.
+    ~5M parameters with default config — suitable for training on consumer GPUs.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.token_emb = nn.Embedding(config.vocab_size, config.dim)
+        self.dropout = nn.Dropout(config.dropout)
+
+        self.layers = nn.ModuleList([
+            TransformerBlock(
+                config.dim, config.n_heads, config.n_kv_heads,
+                config.hidden_dim, config.dropout,
+            )
+            for _ in range(config.n_layers)
+        ])
+
+        self.norm = RMSNorm(config.dim)
+        self.output = nn.Linear(config.dim, config.vocab_size, bias=False)
+
+        # Weight tying — reduces params and improves generalization
+        self.token_emb.weight = self.output.weight
+
+        # Precompute RoPE frequencies
+        head_dim = config.dim // config.n_heads
+        freqs = self._precompute_freqs(head_dim, config.max_seq_len, config.rope_theta)
+        self.register_buffer("freqs", freqs, persistent=False)
+
+        self.grad_checkpoint = False
+        self._init_weights()
+
+    def _precompute_freqs(self, dim: int, max_seq_len: int, theta: float) -> torch.Tensor:
+        freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
+        t = torch.arange(max_seq_len, dtype=torch.float32)
+        return torch.outer(t, freqs)
+
+    def _init_weights(self):
+        """Xavier-style initialization for stable training."""
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.normal_(module.weight, mean=0.0, std=0.02)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+            elif isinstance(module, nn.Embedding):
+                nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        targets: Optional[torch.Tensor] = None,
+        use_cache: bool = False,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        bs, seq_len = input_ids.shape
+        h = self.dropout(self.token_emb(input_ids))
+
+        freqs = self.freqs[:seq_len].to(h.device)
+
+        for layer in self.layers:
+            if self.grad_checkpoint and self.training:
+                h = torch.utils.checkpoint.checkpoint(
+                    layer, h, freqs, None, use_cache, use_reentrant=False
+                )
+            else:
+                h = layer(h, freqs, use_cache=use_cache)
+
+        h = self.norm(h)
+        logits = self.output(h)
+
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(
+                logits.view(-1, logits.size(-1)),
+                targets.view(-1),
+                ignore_index=0,  # Ignore padding
+            )
+
+        return logits, loss
+
+    def reset_caches(self):
+        for layer in self.layers:
+            layer.attention.reset_cache()
+
+    def count_parameters(self) -> int:
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+
+    @classmethod
+    def from_config(cls, model_config) -> "MusicTransformer":
+        return cls(model_config)

src/s05_trainer.py

@@ -0,0 +1,235 @@
+"""
+Training loop with:
+- Mixed precision (AMP) for memory savings
+- Gradient checkpointing for OOM prevention
+- Cosine annealing with warmup
+- Early stopping
+- Gradient accumulation for larger effective batch
+- TensorBoard logging
+"""
+import logging
+import math
+import time
+from pathlib import Path
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
+
+from src.s01_config import TrainConfig, PathConfig, get_device
+from src.s04_model import MusicTransformer
+
+logger = logging.getLogger(__name__)
+
+
+class CosineWarmupScheduler:
+    """Cosine annealing LR with linear warmup."""
+
+    def __init__(self, optimizer, warmup_steps: int, total_steps: int, min_lr: float = 1e-6):
+        self.optimizer = optimizer
+        self.warmup_steps = warmup_steps
+        self.total_steps = total_steps
+        self.min_lr = min_lr
+        self.base_lrs = [pg["lr"] for pg in optimizer.param_groups]
+        self.step_count = 0
+
+    def step(self):
+        self.step_count += 1
+        for pg, base_lr in zip(self.optimizer.param_groups, self.base_lrs):
+            if self.step_count < self.warmup_steps:
+                lr = base_lr * self.step_count / max(1, self.warmup_steps)
+            else:
+                progress = (self.step_count - self.warmup_steps) / max(
+                    1, self.total_steps - self.warmup_steps
+                )
+                lr = self.min_lr + (base_lr - self.min_lr) * 0.5 * (1 + math.cos(math.pi * progress))
+            pg["lr"] = lr
+
+    def get_lr(self) -> float:
+        return self.optimizer.param_groups[0]["lr"]
+
+
+class Trainer:
+    """Handles the full training pipeline with memory-efficient techniques."""
+
+    def __init__(
+        self,
+        model: MusicTransformer,
+        train_loader: DataLoader,
+        val_loader: DataLoader,
+        train_config: TrainConfig,
+        path_config: PathConfig,
+    ):
+        self.model = model
+        self.train_loader = train_loader
+        self.val_loader = val_loader
+        self.config = train_config
+        self.paths = path_config
+        self.device = get_device()
+
+        # Enable gradient checkpointing
+        if train_config.grad_checkpoint:
+            self.model.grad_checkpoint = True
+            logger.info("Gradient checkpointing ENABLED")
+
+        self.model.to(self.device)
+
+        # Optimizer: AdamW with weight decay (decoupled)
+        self.optimizer = torch.optim.AdamW(
+            self.model.parameters(),
+            lr=train_config.learning_rate,
+            weight_decay=train_config.weight_decay,
+            betas=(0.9, 0.95),
+            fused=torch.cuda.is_available(),  # Fused optimizer on CUDA
+        )
+
+        # LR scheduler
+        total_steps = len(train_loader) * train_config.max_epochs // train_config.grad_accum_steps
+        self.scheduler = CosineWarmupScheduler(
+            self.optimizer, train_config.warmup_steps, total_steps
+        )
+
+        # Mixed precision scaler
+        use_amp = train_config.use_amp and torch.cuda.is_available()
+        self.scaler = torch.amp.GradScaler("cuda", enabled=use_amp)
+        if use_amp:
+            self.amp_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+        else:
+            self.amp_dtype = torch.float32
+        self.use_amp = use_amp
+
+        # TensorBoard
+        self.writer = SummaryWriter(log_dir=str(path_config.log_dir))
+
+        # Tracking
+        self.global_step = 0
+        self.best_val_loss = float("inf")
+        self.patience_counter = 0
+
+    def train(self):
+        """Main training loop."""
+        logger.info(f"Starting training on {self.device}")
+        logger.info(f"Model params: {self.model.count_parameters():,}")
+        logger.info(f"Train batches: {len(self.train_loader)}, Val batches: {len(self.val_loader)}")
+
+        for epoch in range(1, self.config.max_epochs + 1):
+            t0 = time.time()
+            train_loss = self._train_epoch(epoch)
+            val_loss = self._validate()
+            elapsed = time.time() - t0
+
+            logger.info(
+                f"Epoch {epoch}/{self.config.max_epochs} | "
+                f"Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f} | "
+                f"LR: {self.scheduler.get_lr():.2e} | Time: {elapsed:.1f}s"
+            )
+
+            self.writer.add_scalars("loss", {"train": train_loss, "val": val_loss}, epoch)
+            self.writer.add_scalar("lr", self.scheduler.get_lr(), epoch)
+
+            # Early stopping check
+            if val_loss < self.best_val_loss - self.config.min_delta:
+                self.best_val_loss = val_loss
+                self.patience_counter = 0
+                self._save_checkpoint("best.pt", epoch, val_loss)
+                logger.info(f"  New best model saved (val_loss={val_loss:.4f})")
+            else:
+                self.patience_counter += 1
+                if self.patience_counter >= self.config.patience:
+                    logger.info(f"Early stopping at epoch {epoch} (patience={self.config.patience})")
+                    break
+
+            # Periodic checkpoint
+            if epoch % 5 == 0:
+                self._save_checkpoint(f"epoch_{epoch}.pt", epoch, val_loss)
+
+        self.writer.close()
+        logger.info("Training complete!")
+
+    def _train_epoch(self, epoch: int) -> float:
+        self.model.train()
+        total_loss = 0.0
+        n_batches = 0
+        self.optimizer.zero_grad(set_to_none=True)
+
+        for batch_idx, (input_ids, targets) in enumerate(self.train_loader):
+            input_ids = input_ids.to(self.device, non_blocking=True)
+            targets = targets.to(self.device, non_blocking=True)
+
+            # Mixed precision forward
+            with torch.amp.autocast(
+                device_type=self.device.type,
+                dtype=self.amp_dtype,
+                enabled=self.use_amp,
+            ):
+                _, loss = self.model(input_ids, targets)
+                loss = loss / self.config.grad_accum_steps
+
+            # Backward with gradient scaling
+            self.scaler.scale(loss).backward()
+
+            if (batch_idx + 1) % self.config.grad_accum_steps == 0:
+                self.scaler.unscale_(self.optimizer)
+                nn.utils.clip_grad_norm_(self.model.parameters(), self.config.max_grad_norm)
+                self.scaler.step(self.optimizer)
+                self.scaler.update()
+                self.optimizer.zero_grad(set_to_none=True)
+                self.scheduler.step()
+                self.global_step += 1
+
+            total_loss += loss.item() * self.config.grad_accum_steps
+            n_batches += 1
+
+            if (batch_idx + 1) % self.config.log_interval == 0:
+                avg = total_loss / n_batches
+                logger.info(
+                    f"  Epoch {epoch} [{batch_idx+1}/{len(self.train_loader)}] "
+                    f"loss={avg:.4f} lr={self.scheduler.get_lr():.2e}"
+                )
+
+        return total_loss / max(1, n_batches)
+
+    @torch.no_grad()
+    def _validate(self) -> float:
+        self.model.eval()
+        total_loss = 0.0
+        n_batches = 0
+
+        for input_ids, targets in self.val_loader:
+            input_ids = input_ids.to(self.device, non_blocking=True)
+            targets = targets.to(self.device, non_blocking=True)
+
+            with torch.amp.autocast(
+                device_type=self.device.type,
+                dtype=self.amp_dtype,
+                enabled=self.use_amp,
+            ):
+                _, loss = self.model(input_ids, targets)
+
+            total_loss += loss.item()
+            n_batches += 1
+
+        return total_loss / max(1, n_batches)
+
+    def _save_checkpoint(self, name: str, epoch: int, val_loss: float):
+        path = self.paths.checkpoint_dir / name
+        torch.save(
+            {
+                "epoch": epoch,
+                "model_state_dict": self.model.state_dict(),
+                "optimizer_state_dict": self.optimizer.state_dict(),
+                "val_loss": val_loss,
+                "global_step": self.global_step,
+                "config": self.model.config,
+            },
+            path,
+        )
+
+    def load_checkpoint(self, path: Path):
+        ckpt = torch.load(path, map_location=self.device, weights_only=False)
+        self.model.load_state_dict(ckpt["model_state_dict"])
+        self.optimizer.load_state_dict(ckpt["optimizer_state_dict"])
+        self.global_step = ckpt.get("global_step", 0)
+        logger.info(f"Loaded checkpoint: {path} (epoch={ckpt['epoch']}, val_loss={ckpt['val_loss']:.4f})")

src/s06_generator.py

@@ -0,0 +1,144 @@
+"""
+Music generation / inference with:
+- Top-k / Top-p (nucleus) sampling
+- Temperature scaling
+- Repetition penalty
+- KV-cache for efficient autoregressive generation
+"""
+import logging
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+from src.s01_config import GenConfig, get_device
+from src.s02_tokenizer import MusicTokenizer, BOS_TOKEN, EOS_TOKEN
+from src.s04_model import MusicTransformer
+
+logger = logging.getLogger(__name__)
+
+
+def top_k_top_p_filter(logits: torch.Tensor, top_k: int, top_p: float) -> torch.Tensor:
+    """Filter logits using top-k and nucleus (top-p) sampling."""
+    if top_k > 0:
+        top_k = min(top_k, logits.size(-1))
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = float("-inf")
+
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+        indices_to_remove = sorted_indices_to_remove.scatter(
+            dim=-1, index=sorted_indices, src=sorted_indices_to_remove
+        )
+        logits[indices_to_remove] = float("-inf")
+
+    return logits
+
+
+def apply_repetition_penalty(logits: torch.Tensor, past_tokens: list[int], penalty: float):
+    """Penalize tokens that have appeared recently."""
+    if penalty == 1.0 or not past_tokens:
+        return logits
+    # Only penalize last 64 tokens to avoid over-suppression
+    recent = past_tokens[-64:]
+    unique_tokens = set(recent)
+    for token_id in unique_tokens:
+        if logits[0, token_id] > 0:
+            logits[0, token_id] /= penalty
+        else:
+            logits[0, token_id] *= penalty
+    return logits
+
+
+@torch.no_grad()
+def generate(
+    model: MusicTransformer,
+    tokenizer: MusicTokenizer,
+    config: GenConfig,
+    prompt_tokens: list[int] | None = None,
+    device: torch.device | None = None,
+) -> list[int]:
+    """
+    Generate music tokens autoregressively using KV-cache.
+
+    Args:
+        model: Trained MusicTransformer
+        tokenizer: MusicTokenizer instance
+        config: Generation config (temperature, top_k, etc.)
+        prompt_tokens: Optional seed tokens (if None, starts with BOS)
+        device: Target device
+    Returns:
+        List of generated token IDs
+    """
+    if device is None:
+        device = get_device()
+
+    model.eval()
+    model.reset_caches()
+    model.to(device)
+
+    if prompt_tokens is None:
+        prompt_tokens = [BOS_TOKEN]
+
+    generated = list(prompt_tokens)
+
+    # Set seed for reproducibility
+    if config.seed is not None:
+        torch.manual_seed(config.seed)
+
+    # Process prompt in one shot (prefill)
+    input_tensor = torch.tensor([generated], dtype=torch.long, device=device)
+    logits, _ = model(input_tensor, use_cache=True)
+
+    for step in range(config.max_tokens - len(generated)):
+        # Get logits for last position
+        next_logits = logits[:, -1, :] / max(config.temperature, 1e-8)
+
+        # Apply repetition penalty
+        next_logits = apply_repetition_penalty(
+            next_logits, generated, config.repetition_penalty
+        )
+
+        # Top-k / Top-p filtering
+        next_logits = top_k_top_p_filter(next_logits, config.top_k, config.top_p)
+
+        # Sample
+        probs = F.softmax(next_logits, dim=-1)
+        next_token = torch.multinomial(probs, num_samples=1).item()
+
+        generated.append(next_token)
+
+        # Stop on EOS
+        if next_token == EOS_TOKEN:
+            break
+
+        # Next step: single token with KV-cache
+        input_tensor = torch.tensor([[next_token]], dtype=torch.long, device=device)
+        logits, _ = model(input_tensor, use_cache=True)
+
+    model.reset_caches()
+    return generated
+
+
+def generate_midi_file(
+    model: MusicTransformer,
+    tokenizer: MusicTokenizer,
+    config: GenConfig,
+    output_path: Path,
+    prompt_tokens: list[int] | None = None,
+):
+    """Generate a MIDI file and save to disk."""
+    logger.info(f"Generating music (max_tokens={config.max_tokens}, temp={config.temperature})...")
+
+    tokens = generate(model, tokenizer, config, prompt_tokens)
+    logger.info(f"Generated {len(tokens)} tokens")
+
+    midi = tokenizer.tokens_to_midi(tokens)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    midi.write(str(output_path))
+    logger.info(f"Saved MIDI to {output_path}")
+    return tokens

src/s07_utils.py

@@ -0,0 +1,49 @@
+"""
+Utility helpers for logging, memory monitoring, and reproducibility.
+"""
+import gc
+import logging
+import os
+import random
+import sys
+
+import numpy as np
+import torch
+
+
+def setup_logging(level: str = "INFO"):
+    logging.basicConfig(
+        level=getattr(logging, level.upper(), logging.INFO),
+        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+        handlers=[logging.StreamHandler(sys.stdout)],
+    )
+
+
+def set_seed(seed: int = 42):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def log_memory_usage(tag: str = ""):
+    """Log current memory usage — critical for constrained hardware."""
+    prefix = f"[{tag}] " if tag else ""
+    if torch.cuda.is_available():
+        allocated = torch.cuda.memory_allocated() / 1024**2
+        reserved = torch.cuda.memory_reserved() / 1024**2
+        logging.info(f"{prefix}GPU Memory: {allocated:.0f}MB allocated, {reserved:.0f}MB reserved")
+    import psutil
+    proc = psutil.Process(os.getpid())
+    ram = proc.memory_info().rss / 1024**2
+    logging.info(f"{prefix}RAM Usage: {ram:.0f}MB")
+
+
+def clear_memory():
+    """Aggressively free memory — call between major operations."""
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()

src/s08_xml_generator.py

@@ -0,0 +1,101 @@
+"""
+XML Configuration-driven Inference for AriaLM.
+Allows specifying prompt notes, generation length, and sampling settings in an XML file.
+"""
+import logging
+import xml.etree.ElementTree as ET
+from pathlib import Path
+import pretty_midi
+import torch
+
+from src.s01_config import GenConfig, ModelConfig, PathConfig
+from src.s02_tokenizer import MusicTokenizer
+from src.s04_model import MusicTransformer
+from src.s06_generator import generate
+
+logger = logging.getLogger(__name__)
+
+
+def parse_xml_config(xml_path: Path) -> tuple[GenConfig, list[int]]:
+    """
+    Parse generation settings and starting prompt notes from an XML file.
+    """
+    tree = ET.parse(xml_path)
+    root = tree.getroot()
+
+    # 1. Parse Settings
+    settings = root.find("settings")
+    temp = float(settings.find("temperature").text or 0.85)
+    top_k = int(settings.find("top_k").text or 40)
+    top_p = float(settings.find("top_p").text or 0.92)
+    max_tokens = int(settings.find("max_tokens").text or 512)
+    rep_penalty = float(settings.find("repetition_penalty").text or 1.15)
+    seed = int(settings.find("seed").text or 42)
+
+    gen_config = GenConfig(
+        temperature=temp,
+        top_k=top_k,
+        top_p=top_p,
+        max_tokens=max_tokens,
+        repetition_penalty=rep_penalty,
+        seed=seed,
+    )
+
+    # 2. Parse Prompt Notes
+    prompt_node = root.find("prompt")
+    prompt_tokens = []
+    if prompt_node is not None:
+        tokenizer = MusicTokenizer()
+        pm = pretty_midi.PrettyMIDI()
+        instrument = pretty_midi.Instrument(program=0)  # Default piano
+
+        current_time = 0.0
+        for note_el in prompt_node.findall("note"):
+            pitch = int(note_el.attrib["pitch"])
+            velocity = int(note_el.attrib["velocity"])
+            duration = float(note_el.attrib["duration_ms"]) / 1000.0
+            delay = float(note_el.attrib.get("delay_ms", 0)) / 1000.0
+
+            current_time += delay
+            note = pretty_midi.Note(
+                velocity=velocity,
+                pitch=pitch,
+                start=current_time,
+                end=current_time + duration,
+            )
+            instrument.notes.append(note)
+            current_time += duration
+
+        pm.instruments.append(instrument)
+        # Convert prompt MIDI to token IDs
+        prompt_tokens = tokenizer.midi_to_tokens(pm)
+
+    return gen_config, prompt_tokens
+
+
+def generate_from_xml(xml_path: Path, model_path: Path, output_path: Path):
+    """Run generation using parameters specified in an XML file."""
+    # Load config and prompt
+    gen_config, prompt_tokens = parse_xml_config(xml_path)
+
+    # Load Model
+    tokenizer = MusicTokenizer()
+    model_config = ModelConfig(vocab_size=tokenizer.vocab_size)
+    model = MusicTransformer.from_config(model_config)
+    ckpt = torch.load(model_path, map_location="cpu", weights_only=False)
+    model.load_state_dict(ckpt["model_state_dict"])
+    
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    
+    print(f"Generating from XML Config: {xml_path}")
+    print(f"Sampling details: Temp={gen_config.temperature}, Seed={gen_config.seed}, Length={gen_config.max_tokens}")
+
+    # Generate
+    tokens = generate(model, tokenizer, gen_config, prompt_tokens, device=device)
+    
+    # Save output
+    midi = tokenizer.tokens_to_midi(tokens)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    midi.write(str(output_path))
+    print(f"Saved generated MIDI to: {output_path}")

tests/test_pipeline.py

@@ -0,0 +1,268 @@
+"""
+End-to-end tests for the music generation pipeline.
+Tests: tokenizer, model, generation, full pipeline.
+"""
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+import torch
+import numpy as np
+from src.s01_config import ModelConfig, GenConfig
+from src.s02_tokenizer import MusicTokenizer, BOS_TOKEN, EOS_TOKEN, VOCAB_SIZE
+from src.s04_model import MusicTransformer
+from src.s06_generator import generate
+
+
+def test_tokenizer_roundtrip():
+    """Test tokenizer encode/decode produces valid events."""
+    tok = MusicTokenizer()
+    assert tok.vocab_size == VOCAB_SIZE
+    assert tok.bos_id == BOS_TOKEN
+    assert tok.eos_id == EOS_TOKEN
+
+    # Test individual token encoding
+    note_on = tok.note_on_token(60)  # Middle C
+    event = tok.decode_token(note_on)
+    assert event["type"] == "NoteOn"
+    assert event["value"] == 60
+
+    vel = tok.velocity_token(100)
+    event = tok.decode_token(vel)
+    assert event["type"] == "Velocity"
+
+    ts = tok.timeshift_token(500)  # 500ms
+    event = tok.decode_token(ts)
+    assert event["type"] == "TimeShift"
+    assert event["value"] == 500
+
+    print("PASS: test_tokenizer_roundtrip")
+
+
+def test_tokenizer_midi_conversion():
+    """Test MIDI to tokens and back."""
+    try:
+        import pretty_midi
+    except ImportError:
+        print("SKIP: test_tokenizer_midi_conversion (pretty_midi not installed)")
+        return
+
+    tok = MusicTokenizer()
+
+    # Create a simple test MIDI
+    midi = pretty_midi.PrettyMIDI(initial_tempo=120)
+    inst = pretty_midi.Instrument(program=0)
+    # C major chord
+    for pitch in [60, 64, 67]:
+        note = pretty_midi.Note(velocity=80, pitch=pitch, start=0.0, end=1.0)
+        inst.notes.append(note)
+    # Second chord at 1.0s
+    for pitch in [65, 69, 72]:
+        note = pretty_midi.Note(velocity=90, pitch=pitch, start=1.0, end=2.0)
+        inst.notes.append(note)
+    midi.instruments.append(inst)
+
+    # Tokenize
+    tokens = tok.midi_to_tokens(midi)
+    assert len(tokens) > 5
+    assert tokens[0] == BOS_TOKEN
+    assert tokens[-1] == EOS_TOKEN
+
+    # Decode back to MIDI
+    midi_out = tok.tokens_to_midi(tokens)
+    assert len(midi_out.instruments) == 1
+    assert len(midi_out.instruments[0].notes) > 0
+
+    print(f"PASS: test_tokenizer_midi_conversion ({len(tokens)} tokens, "
+          f"{len(midi_out.instruments[0].notes)} notes)")
+
+
+def test_model_forward():
+    """Test model forward pass and loss computation."""
+    config = ModelConfig(
+        vocab_size=VOCAB_SIZE,
+        dim=64,
+        n_layers=2,
+        n_heads=4,
+        n_kv_heads=2,
+        max_seq_len=128,
+        dropout=0.0,
+    )
+    model = MusicTransformer.from_config(config)
+
+    # Check parameter count is reasonable
+    n_params = model.count_parameters()
+    assert n_params > 0
+    print(f"  Model params: {n_params:,}")
+
+    # Forward pass
+    batch_size = 2
+    seq_len = 32
+    input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len))
+    targets = torch.randint(0, config.vocab_size, (batch_size, seq_len))
+
+    logits, loss = model(input_ids, targets)
+    assert logits.shape == (batch_size, seq_len, config.vocab_size)
+    assert loss is not None
+    assert loss.item() > 0
+
+    # Backward pass (check gradients flow)
+    loss.backward()
+    grad_norms = [p.grad.norm().item() for p in model.parameters() if p.grad is not None]
+    assert len(grad_norms) > 0
+    assert all(not np.isnan(g) for g in grad_norms)
+
+    print(f"PASS: test_model_forward (loss={loss.item():.4f})")
+
+
+def test_model_gradient_checkpoint():
+    """Test that gradient checkpointing works and reduces memory."""
+    config = ModelConfig(
+        vocab_size=VOCAB_SIZE,
+        dim=64,
+        n_layers=4,
+        n_heads=4,
+        n_kv_heads=2,
+        max_seq_len=128,
+        dropout=0.0,
+    )
+    model = MusicTransformer.from_config(config)
+    model.grad_checkpoint = True
+
+    input_ids = torch.randint(0, config.vocab_size, (2, 64))
+    targets = torch.randint(0, config.vocab_size, (2, 64))
+
+    logits, loss = model(input_ids, targets)
+    loss.backward()
+
+    assert loss.item() > 0
+    print(f"PASS: test_model_gradient_checkpoint (loss={loss.item():.4f})")
+
+
+def test_generation():
+    """Test autoregressive generation."""
+    config = ModelConfig(
+        vocab_size=VOCAB_SIZE,
+        dim=64,
+        n_layers=2,
+        n_heads=4,
+        n_kv_heads=2,
+        max_seq_len=128,
+        dropout=0.0,
+    )
+    model = MusicTransformer.from_config(config)
+    tokenizer = MusicTokenizer()
+
+    gen_config = GenConfig(
+        temperature=0.8,
+        top_k=20,
+        top_p=0.9,
+        max_tokens=50,
+        repetition_penalty=1.1,
+        seed=42,
+    )
+
+    tokens = generate(model, tokenizer, gen_config, device=torch.device("cpu"))
+    assert len(tokens) > 1
+    assert tokens[0] == BOS_TOKEN
+
+    print(f"PASS: test_generation ({len(tokens)} tokens generated)")
+
+
+def test_generation_to_midi():
+    """Test full pipeline: generate tokens → convert to MIDI."""
+    try:
+        import pretty_midi
+    except ImportError:
+        print("SKIP: test_generation_to_midi (pretty_midi not installed)")
+        return
+
+    config = ModelConfig(
+        vocab_size=VOCAB_SIZE,
+        dim=64,
+        n_layers=2,
+        n_heads=4,
+        n_kv_heads=2,
+        max_seq_len=128,
+        dropout=0.0,
+    )
+    model = MusicTransformer.from_config(config)
+    tokenizer = MusicTokenizer()
+
+    gen_config = GenConfig(
+        temperature=1.0,
+        top_k=50,
+        top_p=0.95,
+        max_tokens=100,
+        repetition_penalty=1.1,
+        seed=123,
+    )
+
+    tokens = generate(model, tokenizer, gen_config, device=torch.device("cpu"))
+    midi = tokenizer.tokens_to_midi(tokens)
+
+    assert midi is not None
+    assert len(midi.instruments) == 1
+
+    print(f"PASS: test_generation_to_midi ({len(tokens)} tokens → "
+          f"{len(midi.instruments[0].notes)} notes)")
+
+
+def test_dataset_creation():
+    """Test MidiTokenDataset with synthetic data."""
+    from src.s03_dataset import MidiTokenDataset
+
+    # Synthetic token sequences
+    sequences = [
+        [BOS_TOKEN] + list(np.random.randint(4, VOCAB_SIZE, size=100)) + [EOS_TOKEN]
+        for _ in range(20)
+    ]
+
+    ds = MidiTokenDataset(sequences, max_seq_len=64, pad_id=0)
+    assert len(ds) == 20
+
+    input_ids, targets = ds[0]
+    assert input_ids.shape == (64,)
+    assert targets.shape == (64,)
+    assert input_ids.dtype == torch.long
+
+    print(f"PASS: test_dataset_creation ({len(ds)} sequences)")
+
+
+if __name__ == "__main__":
+    print("=" * 60)
+    print("MUSIC GENERATION LLM — TESTS")
+    print("=" * 60)
+
+    tests = [
+        test_tokenizer_roundtrip,
+        test_tokenizer_midi_conversion,
+        test_model_forward,
+        test_model_gradient_checkpoint,
+        test_generation,
+        test_generation_to_midi,
+        test_dataset_creation,
+    ]
+
+    passed = 0
+    failed = 0
+    skipped = 0
+
+    for test in tests:
+        try:
+            test()
+            passed += 1
+        except Exception as e:
+            if "SKIP" in str(e):
+                skipped += 1
+            else:
+                print(f"FAIL: {test.__name__}: {e}")
+                import traceback
+                traceback.print_exc()
+                failed += 1
+
+    print("=" * 60)
+    print(f"Results: {passed} passed, {failed} failed, {skipped} skipped")
+    print("=" * 60)
+    sys.exit(1 if failed > 0 else 0)