File size: 5,616 Bytes
30e9297 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | """
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()
|