"""End-to-end Whisper encoder regression for Ladder S10/S9. Use under the local fastgen environment. No Whisper decoder is instantiated. """ from __future__ import annotations import argparse from contextlib import nullcontext from functools import partial import io import json import math import os from pathlib import Path from typing import Iterator import numpy as np import soundfile as sf from scipy.signal import resample_poly from safetensors import safe_open import torch import transformers from torch import nn from torch.nn import functional as F from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader, Dataset, Sampler from transformers import WhisperConfig, WhisperFeatureExtractor from transformers.models.whisper.modeling_whisper import WhisperEncoder from timbre_targets import (OrangeTimbreTeacher, OrangeIdentityTeacher, TimbreChunkWriter, TimbreIndex, TimbreTargetCache) ROOT = Path('/e/scratch/reformo/schuhmann1_moss/whisper_score_regression') RUNS = Path('/e/project1/reformo/schuhmann1/whisper_score_regression/runs') MODEL_CACHE = Path('/e/scratch/reformo/schuhmann1_moss/hf_runtime_whisper_regression/hub') _EXTRACTOR = None def snapshot(name: str) -> Path: matches = list((MODEL_CACHE / f'models--openai--whisper-{name}' / 'snapshots').glob('*/config.json')) assert len(matches) == 1, matches return matches[0].parent class AudioRows(Dataset): def __init__(self, arrays: list[dict[str, np.ndarray]], paths: list[str], mean: np.ndarray, std: np.ndarray): self.columns = {key: np.concatenate([part[key] for part in arrays]) for key in arrays[0]} self.paths = paths self.mean, self.std = mean, std self.open_files: dict[int, io.BufferedReader] = {} self.durations = self.columns['dur'] def __len__(self) -> int: return len(self.durations) def __getitem__(self, i: int): c = self.columns path_id = int(c['path'][i]) if path_id not in self.open_files: self.open_files[path_id] = open(self.paths[path_id], 'rb', buffering=0) file = self.open_files[path_id] file.seek(int(c['offset'][i])) raw = file.read(int(c['size'][i])) assert len(raw) == int(c['size'][i]), (i, path_id) audio, sr = sf.read(io.BytesIO(raw), dtype='float32', always_2d=False) if audio.ndim == 2: audio = audio.mean(axis=1) if sr != 16000: common = math.gcd(sr, 16000) audio = resample_poly(audio, 16000 // common, sr // common).astype(np.float32) audio = np.nan_to_num(audio[:480000], copy=False) if len(audio) < 1600: audio = np.pad(audio, (0, 1600 - len(audio))) y = c['y'][i] mask = np.isfinite(y) target = np.where(mask, (y - self.mean) / self.std, 0).astype(np.float32) return audio, target, mask, i def collate(batch, model_path: str, include_audio: bool = False): global _EXTRACTOR if _EXTRACTOR is None: _EXTRACTOR = WhisperFeatureExtractor.from_pretrained(model_path, local_files_only=True) audio, targets, masks, indices = zip(*batch) feat = _EXTRACTOR(list(audio), sampling_rate=16000, padding='longest', truncation=True, max_length=480000, return_attention_mask=True, return_tensors='np') mel = feat['input_features'].astype(np.float32) attention = feat['attention_mask'].astype(np.int64) # The second Whisper convolution has stride two. if mel.shape[-1] % 2: mel = np.pad(mel, ((0, 0), (0, 0), (0, 1))) attention = np.pad(attention, ((0, 0), (0, 1))) result = (torch.from_numpy(mel), torch.from_numpy(attention), torch.from_numpy(np.stack(targets)), torch.from_numpy(np.stack(masks)), indices) return result + (audio,) if include_audio else result def init_worker(_worker_id: int): # OpenMP/ACL was initialized before forking in the parent after CUDA setup. # Spawned workers and one CPU thread avoid inherited OpenMP barriers. torch.set_num_threads(1) class BucketBatches(Sampler[list[int]]): def __init__(self, duration: np.ndarray, batch: int, world: int, rank: int, seed: int): self.duration, self.batch, self.world, self.rank, self.seed = duration, batch, world, rank, seed self.epoch = 0 def __len__(self) -> int: return len(self.duration) // (self.batch * self.world) def __iter__(self) -> Iterator[list[int]]: rng = np.random.default_rng(self.seed + self.epoch) order = rng.permutation(len(self.duration)) width = self.batch * self.world batches = [] for start in range(0, len(order), width * 32): pool = order[start:start + width * 32] pool = pool[np.argsort(self.duration[pool], kind='stable')] for j in range(0, len(pool) - width + 1, width): batches.append(pool[j:j + width]) rng.shuffle(batches) for group in batches: yield group[self.rank * self.batch:(self.rank + 1) * self.batch].tolist() class EncoderRegressor(nn.Module): def __init__(self, model_path: Path, outputs: int, timbre: bool = False, identity: bool = False): super().__init__() config = WhisperConfig.from_pretrained(model_path, local_files_only=True) self.encoder = WhisperEncoder(config) with safe_open(model_path / 'model.safetensors', framework='pt', device='cpu') as file: prefix = 'model.encoder.' state = {key[len(prefix):]: file.get_tensor(key) for key in file.keys() if key.startswith(prefix)} missing, unexpected = self.encoder.load_state_dict(state, strict=True) assert not missing and not unexpected, (missing, unexpected) self.encoder.embed_positions.weight.requires_grad_(True) dim = config.d_model self.head = nn.Sequential(nn.LayerNorm(dim * 4), nn.Linear(dim * 4, 512), nn.GELU(), nn.Dropout(0.1), nn.Linear(512, 256), nn.GELU(), nn.Dropout(0.1), nn.Linear(256, outputs)) self.timbre_head = (nn.Sequential(nn.LayerNorm(dim * 4), nn.Linear(dim * 4, 512), nn.GELU(), nn.Dropout(0.1), nn.Linear(512, 128)) if timbre else None) self.identity_head = (nn.Sequential(nn.LayerNorm(dim * 4), nn.Linear(dim * 4, 512), nn.GELU(), nn.Dropout(0.1), nn.Linear(512, 250)) if identity else None) def forward(self, mel: torch.Tensor, frame_mask: torch.Tensor) -> torch.Tensor: # Identical pretrained encoder blocks, with positional embeddings sliced to the # longest item in this duration bucket. This avoids 30 s padding for every clip. enc = self.encoder h = F.gelu(enc.conv1(mel)) h = F.gelu(enc.conv2(h)).transpose(1, 2) length = h.shape[1] assert length <= enc.max_source_positions h = h + enc.embed_positions(torch.arange(length, device=h.device)) valid = torch.arange(length, device=h.device)[None, :] < ((frame_mask.sum(1) + 1) // 2)[:, None] # Whisper's stock forward ignores masks; here padded frames cannot act as keys. attention = torch.zeros((len(h), 1, 1, length), device=h.device, dtype=h.dtype) attention.masked_fill_(~valid[:, None, None, :], torch.finfo(h.dtype).min) for layer in enc.layers: h = layer(h, attention) h = enc.layer_norm(h).float() mask = valid[:, :, None] count = mask.sum(1).clamp_min(1) mean = (h * mask).sum(1) / count variance = ((h - mean[:, None, :]).square() * mask).sum(1) / count minimum = h.masked_fill(~mask, torch.inf).amin(1) maximum = h.masked_fill(~mask, -torch.inf).amax(1) pooled = torch.cat([mean, minimum, maximum, variance.clamp_min(1e-8).sqrt()], dim=-1) scores = self.head(pooled) if self.timbre_head is None and self.identity_head is None: return scores vectors = {} if self.timbre_head is not None: vectors['timbre'] = F.normalize(self.timbre_head(pooled).float(), dim=-1) if self.identity_head is not None: vectors['identity'] = F.normalize(self.identity_head(pooled).float(), dim=-1) return scores, vectors def regression_loss(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: element = F.huber_loss(pred.float(), target.float(), reduction='none', delta=1.0) return (element * mask).sum() / mask.sum().clamp_min(1) def embedding_loss(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: # Unit vectors: cosine is the primary metric. A small Huber term anchors # individual coordinates without letting dimensions swamp the score loss. target = F.normalize(target.float(), dim=-1) return (1 - (pred.float() * target).sum(-1)).mean() + 0.1 * F.huber_loss( pred.float(), target, delta=0.1) def embedding_targets(caches: dict[str, TimbreTargetCache], columns: dict, indices, audio, device): uids = [columns['uid'][i].decode('utf-8') for i in indices] return {name: torch.from_numpy(cache.get_or_compute(uids, audio).copy()).to( device, dtype=torch.float32) for name, cache in caches.items()} @torch.no_grad() def evaluate(model, loader, mean: torch.Tensor, std: torch.Tensor, device, world: int, max_batches: int = 0, embedding_caches: dict[str, TimbreTargetCache] | None = None, columns: dict | None = None): model.eval() n = len(mean) sums = torch.zeros((9, n), dtype=torch.float64, device=device) embedding_caches = embedding_caches or {} embedding_sums = {name: torch.zeros(3, dtype=torch.float64, device=device) for name in embedding_caches} for step, batch in enumerate(loader): if max_batches and step >= max_batches: break mel, frame_mask, target, mask, indices = batch[:5] mel, frame_mask = mel.to(device), frame_mask.to(device) target, mask = target.to(device), mask.to(device) if embedding_caches: assert columns is not None target_vectors = embedding_targets(embedding_caches, columns, indices, batch[5], device) with torch.autocast('cuda', dtype=torch.bfloat16): output = model(mel, frame_mask) if embedding_caches: pred, pred_vectors = output for name, target_vector in target_vectors.items(): pred_vector = pred_vectors[name] stats = embedding_sums[name] stats[0] += len(target_vector) stats[1] += (pred_vector * F.normalize(target_vector, dim=-1)).sum().double() stats[2] += F.huber_loss(pred_vector, target_vector, delta=0.1, reduction='sum').double() / target_vector.shape[1] else: pred = output pred, target, mask = pred.double(), target.double(), mask.double() y = target * std + mean p = pred * std + mean sums[0] += mask.sum(0) sums[1] += (F.huber_loss(pred, target, reduction='none', delta=1.0) * mask).sum(0) sums[2] += ((p - y).abs() * mask).sum(0) sums[3] += ((p - y).square() * mask).sum(0) sums[4] += (y * mask).sum(0) sums[5] += (y.square() * mask).sum(0) sums[6] += (p * mask).sum(0) sums[7] += (p.square() * mask).sum(0) sums[8] += (p * y * mask).sum(0) if world > 1: torch.distributed.all_reduce(sums) for stats in embedding_sums.values(): torch.distributed.all_reduce(stats) count = sums[0].clamp_min(1) mae = sums[2] / count huber = sums[1] / count y_var = (sums[5] - sums[4].square() / count).clamp_min(1e-10) p_var = (sums[7] - sums[6].square() / count).clamp_min(1e-10) corr = (sums[8] - sums[4] * sums[6] / count) / (y_var * p_var).sqrt() r2 = 1 - sums[3] / y_var result = {'count': sums[0].int().tolist(), 'mae': mae.tolist(), 'huber_z': huber.tolist(), 'pearson': corr.tolist(), 'r2': r2.tolist()} for name, a, b in [('emotion', 0, 40), ('voicenet', 40, 97), ('other', 97, 100)]: result[name] = {'mae': float(mae[a:b].mean()), 'huber_z': float(huber[a:b].mean()), 'pearson': float(corr[a:b].mean()), 'r2': float(r2[a:b].mean())} result['huber_z_macro'] = float(huber.mean()) for name, stats in embedding_sums.items(): result[f'{name}_count'] = int(stats[0]) result[f'{name}_cosine'] = float(stats[1] / stats[0].clamp_min(1)) result[f'{name}_huber'] = float(stats[2] / stats[0].clamp_min(1)) return result def main(): p = argparse.ArgumentParser() p.add_argument('--model', choices=['tiny', 'base'], required=True) p.add_argument('--stage', choices=['S10', 'S9'], required=True) p.add_argument('--epochs', type=int, default=3) p.add_argument('--batch-per-gpu', type=int, default=16) p.add_argument('--accum', type=int, default=2) p.add_argument('--workers', type=int, default=2) p.add_argument('--encoder-lr', type=float, default=2e-5) p.add_argument('--head-lr', type=float, default=3e-4) p.add_argument('--timbre', action='store_true', help='Train the 128-D Orange TBR head') p.add_argument('--timbre-weight', type=float, default=0.2) p.add_argument('--timbre-index', type=Path, help='Sorted index of precomputed Emolia embeddings') p.add_argument('--timbre-generated-index', type=Path, help='Index of earlier generated chunks') p.add_argument('--timbre-cache-dir', type=Path, default=ROOT / 'timbre' / 'chunks') p.add_argument('--identity', action='store_true', help='Train the 250-D Orange ID head') p.add_argument('--identity-weight', type=float, default=0.2) p.add_argument('--identity-index', type=Path, help='Sorted index of precomputed KartoffelPhon embeddings') p.add_argument('--identity-generated-index', type=Path, help='Index of earlier generated ID chunks') p.add_argument('--identity-cache-dir', type=Path, default=ROOT / 'identity' / 'chunks') p.add_argument('--max-train-batches', type=int, default=0) p.add_argument('--max-val-batches', type=int, default=0) p.add_argument('--out', type=Path) args = p.parse_args() assert torch.cuda.is_available(), 'Run on a Booster GPU node' world = int(os.environ.get('WORLD_SIZE', os.environ.get('SLURM_NTASKS', '1'))) rank = int(os.environ.get('RANK', os.environ.get('SLURM_PROCID', '0'))) if world > 1: os.environ.setdefault('RANK', str(rank)) os.environ.setdefault('WORLD_SIZE', str(world)) torch.distributed.init_process_group('nccl') torch.cuda.set_device(0) torch.set_num_threads(2) torch.manual_seed(20260928) device = torch.device('cuda:0') meta = json.loads((ROOT / 'manifest.json').read_text()) mean = torch.tensor(meta['mean'], dtype=torch.float64, device=device) std = torch.tensor(meta['std'], dtype=torch.float64, device=device) def load(name): with np.load(ROOT / f'{name}.npz') as file: return {key: file[key] for key in file.files} train_parts = [load('t0')] + ([load('t1')] if args.stage == 'S9' else []) train_data = AudioRows(train_parts, meta['paths'], np.asarray(meta['mean']), np.asarray(meta['std'])) val_data = AudioRows([load('val')], meta['paths'], np.asarray(meta['mean']), np.asarray(meta['std'])) path = snapshot(args.model) sampler = BucketBatches(train_data.durations, args.batch_per_gpu, world, rank, meta['seed']) collate_fn = partial(collate, model_path=str(path), include_audio=args.timbre or args.identity) loader_options = {'num_workers': args.workers, 'pin_memory': True, 'persistent_workers': args.workers > 0, 'collate_fn': collate_fn} if args.workers: loader_options.update({'multiprocessing_context': 'spawn', 'worker_init_fn': init_worker}) train_loader = DataLoader(train_data, batch_sampler=sampler, **loader_options) val_indices = np.arange(rank, len(val_data), world) val_indices = val_indices[np.argsort(val_data.durations[val_indices])] val_batches = [val_indices[i:i + args.batch_per_gpu].tolist() for i in range(0, len(val_indices), args.batch_per_gpu)] val_loader = DataLoader(val_data, batch_sampler=val_batches, **loader_options) core = EncoderRegressor(path, len(meta['targets']), timbre=args.timbre, identity=args.identity).to(device) model = DDP(core, device_ids=[0], broadcast_buffers=False) if world > 1 else core parameter_groups = [ {'params': core.encoder.parameters(), 'lr': args.encoder_lr}, {'params': core.head.parameters(), 'lr': args.head_lr}] if args.timbre: parameter_groups.append({'params': core.timbre_head.parameters(), 'lr': args.head_lr}) if args.identity: parameter_groups.append({'params': core.identity_head.parameters(), 'lr': args.head_lr}) optimizer = torch.optim.AdamW(parameter_groups, weight_decay=0.01) embedding_caches = {} if args.timbre: assert args.timbre_index is not None and (args.timbre_index / 'manifest.json').exists(), \ 'Build the Emolia timbre index first: python timbre_targets.py index-emolia --output PATH' teacher = OrangeTimbreTeacher().to(device).eval() for parameter in teacher.parameters(): parameter.requires_grad_(False) precomputed = TimbreIndex(args.timbre_index, expected_dim=128) generated = TimbreIndex(args.timbre_generated_index, expected_dim=128) if args.timbre_generated_index else None writer = TimbreChunkWriter(args.timbre_cache_dir, rank, flush_rows=4096) embedding_caches['timbre'] = TimbreTargetCache(precomputed, generated, teacher, writer, device) if args.identity: assert args.identity_index is not None and (args.identity_index / 'manifest.json').exists(), \ 'Build the Orange ID index first: python timbre_targets.py index-identity --output PATH' teacher = OrangeIdentityTeacher().to(device).eval() for parameter in teacher.parameters(): parameter.requires_grad_(False) precomputed = TimbreIndex(args.identity_index, expected_dim=250) generated = TimbreIndex(args.identity_generated_index, expected_dim=250) if args.identity_generated_index else None writer = TimbreChunkWriter(args.identity_cache_dir, rank, flush_rows=4096, dim=250, model='Orange/Speaker-wavLM-id', revision='abbb3c7b8d220ceebc33d9bd3bb5aedb580342c') embedding_caches['identity'] = TimbreTargetCache(precomputed, generated, teacher, writer, device) batches = min(len(sampler), args.max_train_batches or len(sampler)) total_steps = args.epochs * math.ceil(batches / args.accum) warmup = max(1, int(total_steps * 0.05)) def schedule(step): if step < warmup: return (step + 1) / warmup progress = (step - warmup) / max(1, total_steps - warmup) return 0.05 + 0.95 * 0.5 * (1 + math.cos(math.pi * progress)) scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, schedule) suffix = '_timbre_identity' if args.timbre and args.identity else \ '_timbre' if args.timbre else '_identity' if args.identity else '' out = args.out or RUNS / f'{args.model}_{args.stage.lower()}{suffix}' out.mkdir(parents=True, exist_ok=True) if rank == 0: (out / 'config.json').write_text(json.dumps(vars(args) | { 'out': str(out), 'world': world, 'train_rows': len(train_data), 'val_rows': len(val_data), 'model_snapshot': str(path), 'torch_version': torch.__version__, 'transformers_version': transformers.__version__, }, indent=2, default=str) + '\n') print('CONFIG', args.model, args.stage, 'train', len(train_data), 'val', len(val_data), 'world', world, 'batches', batches, flush=True) best = math.inf updates = 0 for epoch in range(args.epochs): sampler.epoch = epoch model.train() optimizer.zero_grad(set_to_none=True) running = 0.0 for step, batch in enumerate(train_loader): if step >= batches: break mel, frame_mask, target, mask, indices = batch[:5] mel, frame_mask = mel.to(device, non_blocking=True), frame_mask.to(device, non_blocking=True) target, mask = target.to(device, non_blocking=True), mask.to(device, non_blocking=True) if embedding_caches: target_vectors = embedding_targets(embedding_caches, train_data.columns, indices, batch[5], device) sync = ((step + 1) % args.accum == 0) or (step + 1 == batches) context = nullcontext() if sync or world == 1 else model.no_sync() with context: with torch.autocast('cuda', dtype=torch.bfloat16): output = model(mel, frame_mask) if embedding_caches: pred, pred_vectors = output loss = regression_loss(pred, target, mask) for name, target_vector in target_vectors.items(): weight = args.timbre_weight if name == 'timbre' else args.identity_weight loss = loss + weight * embedding_loss(pred_vectors[name], target_vector) else: loss = regression_loss(output, target, mask) group_start = (step // args.accum) * args.accum group_size = min(args.accum, batches - group_start) (loss / group_size).backward() running += float(loss.detach()) if sync: nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step() optimizer.zero_grad(set_to_none=True) updates += 1 if rank == 0 and (updates % 50 == 0 or step + 1 == batches): print('STEP', epoch + 1, updates, step + 1, 'loss', running / (step + 1), flush=True) metrics = evaluate(core, val_loader, mean, std, device, world, args.max_val_batches, embedding_caches, val_data.columns) for name, cache in embedding_caches.items(): cache.writer.flush() count = torch.tensor(cache.computed, dtype=torch.long, device=device) if world > 1: torch.distributed.all_reduce(count) metrics[f'{name}_new_saved'] = int(count) if rank == 0: selection = metrics['huber_z_macro'] for name in embedding_caches: weight = args.timbre_weight if name == 'timbre' else args.identity_weight selection += weight * (1 - metrics[f'{name}_cosine']) metrics.update({'epoch': epoch + 1, 'train_loss': running / batches, 'selection_loss': selection, 'model': args.model, 'stage': args.stage, 'updates': updates}) if not embedding_caches: metrics['train_huber_z'] = running / batches with (out / 'metrics.jsonl').open('a') as file: file.write(json.dumps(metrics) + '\n') print('EPOCH', epoch + 1, 'train', metrics['train_loss'], 'val', metrics['huber_z_macro'], 'emo', metrics['emotion'], 'vn', metrics['voicenet'], flush=True) if selection < best: best = selection torch.save({'model': core.state_dict(), 'epoch': epoch + 1, 'targets': meta['targets'], 'mean': meta['mean'], 'std': meta['std']}, out / 'best.pt') if world > 1: torch.distributed.barrier() if world > 1: torch.distributed.destroy_process_group() if rank == 0 and not args.max_train_batches and not embedding_caches: from report import main as write_comparison write_comparison() if __name__ == '__main__': main()