"""Train compound MIDI-text contrastive model from a CompoundGPT checkpoint.""" from __future__ import annotations import argparse import csv import json import math import random import sys import time from pathlib import Path from typing import Any, Dict, List, Optional, Sequence, Tuple import pretty_midi import torch from torch.optim import AdamW from torch.utils.data import DataLoader, Dataset _SCRIPT_DIR = Path(__file__).resolve().parent _ROOT = _SCRIPT_DIR.parent if str(_SCRIPT_DIR) not in sys.path: sys.path.insert(0, str(_SCRIPT_DIR)) from compound import SENTINELS, STEP_PAD, encode_compound # noqa: E402 from compound_model import ( # noqa: E402 CompoundGPT, CompoundGPTConfig, default_compound_config, ) from contrastive_model import CompoundMidiTextContrastiveModel # noqa: E402 def _pick_device() -> torch.device: if torch.cuda.is_available(): return torch.device("cuda") mps = getattr(torch.backends, "mps", None) if mps is not None and mps.is_available(): return torch.device("mps") return torch.device("cpu") def _load_jsonl_records(path: Path) -> List[Dict[str, Any]]: records: List[Dict[str, Any]] = [] with open(path) as f: for line in f: line = line.strip() if line: records.append(json.loads(line)) return records def _caption_from_record(record: Dict[str, Any]) -> str: llm_caption = str(record.get("caption", "")).strip() if llm_caption: return llm_caption return str(record.get("caption_template", "")).strip() def _path_from_record(record: Dict[str, Any]) -> str: for k in ("path", "midi_path", "midi"): value = str(record.get(k, "")).strip() if value: return value return "" class CompoundCaptionDataset(Dataset): """Dataset yielding padded compound steps + captions.""" def __init__( self, records: Sequence[Dict[str, Any]], is_train: bool, max_seq_len: int, seed: int, ) -> None: self.records = list(records) self.is_train = is_train self.max_seq_len = max_seq_len self._rng = random.Random(seed) self._cache: Dict[int, Optional[List[List[int]]]] = {} self._valid_indices: set[int] = set() def __len__(self) -> int: return len(self.records) def _encode_idx(self, idx: int) -> Optional[List[List[int]]]: if idx in self._cache: return self._cache[idx] rec = self.records[idx] path = _path_from_record(rec) if not path: self._cache[idx] = None return None try: pm = pretty_midi.PrettyMIDI(path) steps = encode_compound(pm) except Exception: self._cache[idx] = None return None if not steps: self._cache[idx] = None return None self._cache[idx] = steps self._valid_indices.add(idx) return steps def _crop(self, steps: List[List[int]]) -> List[List[int]]: if self.is_train: max_start = len(steps) - self.max_seq_len start = 0 if max_start <= 0 else self._rng.randint(0, max_start) return steps[start : start + self.max_seq_len] return steps[: self.max_seq_len] def _pad(self, steps: List[List[int]]) -> List[List[int]]: if len(steps) >= self.max_seq_len: return steps pad_step = list(SENTINELS) pad_step[0] = STEP_PAD return steps + [pad_step] * (self.max_seq_len - len(steps)) def __getitem__(self, idx: int) -> Dict[str, Any]: steps = self._encode_idx(idx) rec = self.records[idx] if steps is None: if not self._valid_indices: for j in range(len(self.records)): self._encode_idx(j) if not self._valid_indices: raise RuntimeError("No valid MIDI records for compound contrastive.") repl_idx = self._rng.choice(list(self._valid_indices)) steps = self._encode_idx(repl_idx) rec = self.records[repl_idx] assert steps is not None steps = self._pad(self._crop(steps)) return { "compound_input": torch.tensor(steps, dtype=torch.long), "caption": _caption_from_record(rec), "path": _path_from_record(rec), } def _collate(items: List[Dict[str, Any]]) -> Dict[str, Any]: return { "compound_input": torch.stack([x["compound_input"] for x in items], dim=0), "captions": [x["caption"] for x in items], "paths": [x["path"] for x in items], } def _split_records( records: List[Dict[str, Any]], split_ratio: float, seed: int ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: rng = random.Random(seed) idx = list(range(len(records))) rng.shuffle(idx) n_train = int(len(idx) * split_ratio) train = [records[i] for i in idx[:n_train]] val = [records[i] for i in idx[n_train:]] return train, val def _set_cosine_lrs( optimizer: AdamW, current_step: int, total_steps: int, min_lr_scale: float ) -> None: if total_steps <= 1: mult = 1.0 else: progress = min(1.0, max(0.0, current_step / float(total_steps - 1))) mult = min_lr_scale + (1.0 - min_lr_scale) * 0.5 * ( 1.0 + math.cos(math.pi * progress) ) for group in optimizer.param_groups: base_lr = group.get("initial_lr", group["lr"]) group["lr"] = base_lr * mult @torch.no_grad() def evaluate( model: CompoundMidiTextContrastiveModel, loader: DataLoader, device: torch.device, ) -> Dict[str, float]: model.eval() loss_sum = 0.0 n_batches = 0 midi_chunks: List[torch.Tensor] = [] text_chunks: List[torch.Tensor] = [] for batch in loader: out = model( compound_input=batch["compound_input"].to(device), captions=batch["captions"], ) loss_sum += float(out["loss"].item()) midi_chunks.append(out["midi_embeds"].detach().cpu()) text_chunks.append(out["text_embeds"].detach().cpu()) n_batches += 1 model.train() if n_batches == 0: return {"loss": 0.0, "r1_m2t": 0.0, "r1_t2m": 0.0} midi_all = torch.cat(midi_chunks, dim=0) text_all = torch.cat(text_chunks, dim=0) logits = midi_all @ text_all.t() labels = torch.arange(logits.size(0)) r1_m2t = float((torch.argmax(logits, dim=1) == labels).float().mean().item()) r1_t2m = float((torch.argmax(logits.t(), dim=1) == labels).float().mean().item()) return {"loss": loss_sum / n_batches, "r1_m2t": r1_m2t, "r1_t2m": r1_t2m} def _load_compound_gpt(ckpt_path: Path, device: torch.device) -> CompoundGPT: ckpt = torch.load(ckpt_path, map_location=device, weights_only=True) cfg = default_compound_config() raw_cfg = ckpt.get("config") if isinstance(ckpt, dict) else None if isinstance(raw_cfg, dict): for k in CompoundGPTConfig.__dataclass_fields__.keys(): if k in raw_cfg: setattr(cfg, k, raw_cfg[k]) model = CompoundGPT(cfg).to(device) state = ckpt.get("model_state_dict", ckpt) model.load_state_dict(state, strict=False) model.eval() return model def train(args: argparse.Namespace) -> None: device = _pick_device() print(f"[compound-contrastive] device={device}") torch.manual_seed(args.seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(args.seed) all_records = _load_jsonl_records(Path(args.captions_jsonl)) train_records, val_records = _split_records( all_records, split_ratio=args.split_ratio, seed=args.seed ) print( "[compound-contrastive] records total/train/val=" f"{len(all_records)}/{len(train_records)}/{len(val_records)}" ) train_ds = CompoundCaptionDataset( records=train_records, is_train=True, max_seq_len=args.max_seq_len, seed=args.seed, ) val_ds = CompoundCaptionDataset( records=val_records, is_train=False, max_seq_len=args.max_seq_len, seed=args.seed + 1, ) train_loader = DataLoader( train_ds, batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=args.num_workers, pin_memory=torch.cuda.is_available(), collate_fn=_collate, ) val_loader = DataLoader( val_ds, batch_size=args.batch_size, shuffle=False, drop_last=False, num_workers=args.num_workers, pin_memory=torch.cuda.is_available(), collate_fn=_collate, ) midi_encoder = _load_compound_gpt( ckpt_path=Path(args.compound_checkpoint), device=device, ) model = CompoundMidiTextContrastiveModel( midi_compound_gpt=midi_encoder, text_model_name=args.text_model, embed_dim=args.embed_dim, init_temperature=args.init_temperature, min_temperature=args.min_temperature, max_temperature=args.max_temperature, device=device, ) model.train() optimizer = AdamW( [ { "params": model.midi_projection.parameters(), "lr": args.proj_lr, "initial_lr": args.proj_lr, "weight_decay": args.proj_weight_decay, }, { "params": model.text_projection.parameters(), "lr": args.proj_lr, "initial_lr": args.proj_lr, "weight_decay": args.proj_weight_decay, }, { "params": [model.log_temperature], "lr": args.temp_lr, "initial_lr": args.temp_lr, "weight_decay": 0.0, }, ] ) steps_per_epoch = len(train_loader) total_steps = max(1, args.epochs * steps_per_epoch) global_step = 0 best_val = float("inf") text_group_added = False start_epoch = 1 if args.resume_checkpoint: ckpt = torch.load( args.resume_checkpoint, map_location=device, weights_only=False, ) model.load_state_dict(ckpt["model_state_dict"]) if int(ckpt["epoch"]) >= args.unfreeze_text_epoch: model.unfreeze_text_encoder() optimizer.add_param_group( { "params": model.text_encoder.parameters(), "lr": args.text_lr, "initial_lr": args.text_lr, "weight_decay": args.text_weight_decay, } ) text_group_added = True optimizer.load_state_dict(ckpt["optimizer_state_dict"]) start_epoch = int(ckpt["epoch"]) + 1 global_step = int(ckpt["global_step"]) best_val = float(ckpt["val_loss"]) print( "[compound-contrastive] resumed from " f"epoch={ckpt['epoch']} val_loss={best_val:.4f}" ) results_dir = Path(args.results_dir) ckpt_dir = results_dir / "checkpoints_contrastive_compound" ckpt_dir.mkdir(parents=True, exist_ok=True) results_dir.mkdir(parents=True, exist_ok=True) log_csv = results_dir / "contrastive_compound_training_log.csv" if not log_csv.exists(): with open(log_csv, "w", newline="") as f: csv.DictWriter( f, fieldnames=[ "epoch", "step", "stage", "train_loss", "train_acc_m2t", "train_acc_t2m", "val_loss", "val_r1_m2t", "val_r1_t2m", "temperature", ], ).writeheader() t0 = time.perf_counter() for epoch in range(start_epoch, args.epochs + 1): if epoch == args.unfreeze_text_epoch and not text_group_added: model.unfreeze_text_encoder() optimizer.add_param_group( { "params": model.text_encoder.parameters(), "lr": args.text_lr, "initial_lr": args.text_lr, "weight_decay": args.text_weight_decay, } ) text_group_added = True print(f"[compound-contrastive] epoch={epoch}: unfroze text encoder.") stage = "A" if epoch < args.unfreeze_text_epoch else "B" train_loss_sum = 0.0 train_acc_m2t_sum = 0.0 train_acc_t2m_sum = 0.0 train_batches = 0 for batch in train_loader: _set_cosine_lrs( optimizer=optimizer, current_step=global_step, total_steps=total_steps, min_lr_scale=args.min_lr_scale, ) optimizer.zero_grad(set_to_none=True) out = model( compound_input=batch["compound_input"].to(device), captions=batch["captions"], ) loss = out["loss"] loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip_norm) optimizer.step() train_loss_sum += float(loss.item()) train_acc_m2t_sum += float(out["acc_midi_to_text"].item()) train_acc_t2m_sum += float(out["acc_text_to_midi"].item()) train_batches += 1 global_step += 1 train_loss = train_loss_sum / max(1, train_batches) train_acc_m2t = train_acc_m2t_sum / max(1, train_batches) train_acc_t2m = train_acc_t2m_sum / max(1, train_batches) val = evaluate(model, val_loader, device) temp_val = float(model.get_temperature().item()) print( f"[compound-contrastive] epoch={epoch}/{args.epochs} stage={stage} " f"train_loss={train_loss:.4f} val_loss={val['loss']:.4f} " f"train_acc=({train_acc_m2t:.3f},{train_acc_t2m:.3f}) " f"val_r1=({val['r1_m2t']:.3f},{val['r1_t2m']:.3f}) " f"temp={temp_val:.4f}" ) with open(log_csv, "a", newline="") as f: csv.DictWriter( f, fieldnames=[ "epoch", "step", "stage", "train_loss", "train_acc_m2t", "train_acc_t2m", "val_loss", "val_r1_m2t", "val_r1_t2m", "temperature", ], ).writerow( { "epoch": epoch, "step": global_step, "stage": stage, "train_loss": f"{train_loss:.6f}", "train_acc_m2t": f"{train_acc_m2t:.6f}", "train_acc_t2m": f"{train_acc_t2m:.6f}", "val_loss": f"{val['loss']:.6f}", "val_r1_m2t": f"{val['r1_m2t']:.6f}", "val_r1_t2m": f"{val['r1_t2m']:.6f}", "temperature": f"{temp_val:.6f}", } ) ckpt = { "model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "epoch": epoch, "global_step": global_step, "val_loss": float(val["loss"]), "args": vars(args), } torch.save(ckpt, ckpt_dir / "clap_compound_latest.pt") if epoch % args.checkpoint_every == 0: torch.save(ckpt, ckpt_dir / f"clap_compound_epoch_{epoch:03d}.pt") if val["loss"] < best_val: best_val = val["loss"] torch.save(ckpt, ckpt_dir / "clap_compound_best.pt") elapsed = time.perf_counter() - t0 print( f"[compound-contrastive] finished in {elapsed/60:.1f} min, " f"best_val={best_val:.4f}" ) def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(description="Train compound MIDI-text contrastive model") p.add_argument( "--captions-jsonl", type=str, default=str(_ROOT / "data" / "midicaps_100.jsonl"), ) p.add_argument( "--compound-checkpoint", type=str, default=str( _ROOT / "results" / "test_compound" / "checkpoints_compound" / "compound_best.pt" ), ) p.add_argument( "--text-model", type=str, default="sentence-transformers/all-MiniLM-L6-v2", ) p.add_argument("--results-dir", type=str, default=str(_ROOT / "results" / "test_compound")) p.add_argument("--epochs", type=int, default=10) p.add_argument("--unfreeze-text-epoch", type=int, default=6) p.add_argument("--batch-size", type=int, default=16) p.add_argument("--max-seq-len", type=int, default=512) p.add_argument("--split-ratio", type=float, default=0.95) p.add_argument("--num-workers", type=int, default=4) p.add_argument("--seed", type=int, default=17) p.add_argument("--embed-dim", type=int, default=256) p.add_argument("--init-temperature", type=float, default=0.07) p.add_argument("--min-temperature", type=float, default=0.01) p.add_argument("--max-temperature", type=float, default=1.0) p.add_argument("--proj-lr", type=float, default=1e-4) p.add_argument("--temp-lr", type=float, default=1e-5) p.add_argument("--text-lr", type=float, default=1e-5) p.add_argument("--proj-weight-decay", type=float, default=0.01) p.add_argument("--text-weight-decay", type=float, default=0.01) p.add_argument("--min-lr-scale", type=float, default=0.01) p.add_argument("--grad-clip-norm", type=float, default=1.0) p.add_argument("--checkpoint-every", type=int, default=5) p.add_argument("--resume-checkpoint", type=str, default="") return p.parse_args() if __name__ == "__main__": train(parse_args())