docs/flow_diagram.drawio

<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>