Release best Gemini-tuned Whisper Base and Small with code, normalization and evaluation
cd9b2d8 verified | """Index the released Ladder TARs without copying audio; make a fixed validation set.""" | |
| from __future__ import annotations | |
| import json | |
| import re | |
| import tarfile | |
| from pathlib import Path | |
| import numpy as np | |
| import pyarrow.parquet as pq | |
| ROOT = Path('/e/scratch/reformo/schuhmann1_moss/publish/scaling_ladder') | |
| OUT = Path('/e/scratch/reformo/schuhmann1_moss/whisper_score_regression') | |
| SEED = 20260928 | |
| def target_names() -> list[str]: | |
| text = (ROOT / 'COLUMNS.md').read_text() | |
| columns = re.findall(r'^\| `([^`]+)`', text, re.M) | |
| emo = [c for c in columns if c.startswith('emo_')] | |
| voice = [c for c in columns if c.startswith('vn_') and c.endswith('_reg')] | |
| assert len(emo) == 40 and len(voice) == 57, (len(emo), len(voice)) | |
| return emo + voice + ['genuineness_0_6', 'blend_0_10', 'R_quality'] | |
| def build() -> None: | |
| OUT.mkdir(parents=True, exist_ok=True) | |
| targets = target_names() | |
| wanted = ['uid', 'dur_s', 'lang', 'origin', 'source'] + targets | |
| paths: list[str] = [] | |
| sources: list[str] = [] | |
| accum: dict[str, dict[str, list[np.ndarray]]] = { | |
| s: {k: [] for k in ['path', 'offset', 'size', 'dur', 'uid', 'cell', 'source', 'y']} | |
| for s in ['t0', 't1', 'heldout'] | |
| } | |
| for stage, pattern in [('t0', '*_t0-*.parquet'), ('t1', '*_t1-*.parquet'), ('heldout', '*_h-*.parquet')]: | |
| for index_path in sorted((ROOT / 'index').rglob(pattern)): | |
| tar_path = ROOT / 'data' / index_path.relative_to(ROOT / 'index').with_suffix('.tar') | |
| assert tar_path.is_file(), tar_path | |
| table = pq.read_table(index_path, columns=wanted) | |
| uids = table['uid'].to_pylist() | |
| n = len(uids) | |
| offsets = np.empty(n, dtype=np.int64) | |
| sizes = np.empty(n, dtype=np.int32) | |
| with tarfile.open(tar_path, 'r:') as archive: | |
| j = 0 | |
| for member in archive: | |
| if not member.name.endswith('.mp3'): | |
| continue | |
| assert j < n and member.name.rsplit('/', 1)[-1][:-4] == uids[j], (tar_path, j) | |
| offsets[j], sizes[j] = member.offset_data, member.size | |
| j += 1 | |
| assert j == n, (tar_path, j, n) | |
| path_id = len(paths) | |
| paths.append(str(tar_path)) | |
| lang = table['lang'].to_pylist() | |
| origin = table['origin'].to_pylist() | |
| cells = np.array([{'de_real': 0, 'de_synth': 1, 'en_real': 2, 'en_synth': 3}[f'{a}_{b}'] | |
| for a, b in zip(lang, origin)], dtype=np.uint8) | |
| source_names = table['source'].to_pylist() | |
| for name in source_names: | |
| if name not in sources: | |
| sources.append(name) | |
| source_id = np.array([sources.index(name) for name in source_names], dtype=np.uint8) | |
| label_cols = [np.asarray(table[name].to_numpy(zero_copy_only=False), dtype=np.float32) | |
| for name in targets] | |
| labels = np.stack(label_cols, axis=1) | |
| chunk = { | |
| 'path': np.full(n, path_id, dtype=np.uint16), | |
| 'offset': offsets, 'size': sizes, | |
| 'dur': np.asarray(table['dur_s'].to_numpy(), dtype=np.float32), | |
| 'uid': np.asarray(uids, dtype='S128'), | |
| 'cell': cells, 'source': source_id, 'y': labels, | |
| } | |
| for name, value in chunk.items(): | |
| accum[stage][name].append(value) | |
| print(stage, tar_path.name, n, flush=True) | |
| data = {stage: {name: np.concatenate(parts) for name, parts in columns.items()} | |
| for stage, columns in accum.items()} | |
| # Shared normalization must not reveal S9-only labels to the S10 run. | |
| train = data['t0']['y'] | |
| with np.errstate(invalid='ignore'): | |
| mean = np.nanmean(train, axis=0).astype(np.float32) | |
| std = np.nanstd(train, axis=0).astype(np.float32) | |
| assert np.all(np.isfinite(mean)) and np.all(np.isfinite(std)) | |
| std = np.maximum(std, 0.25) | |
| # Separate held-out UIDs are a release invariant; fail if it changed. | |
| train_uids = set(np.concatenate([data['t0']['uid'], data['t1']['uid']]).tolist()) | |
| val_uids = set(data['heldout']['uid'].tolist()) | |
| assert len(train_uids) == len(data['t0']['uid']) + len(data['t1']['uid']) | |
| assert not (train_uids & val_uids) | |
| selected = balanced_val(data['heldout'], 2048) | |
| for stage, arrays in data.items(): | |
| if stage == 'heldout': | |
| arrays = {name: value[selected] for name, value in arrays.items()} | |
| stage = 'val' | |
| np.savez(OUT / f'{stage}.npz', **arrays) | |
| meta = { | |
| 'root': str(ROOT), 'seed': SEED, 'targets': targets, 'paths': paths, | |
| 'sources': sources, 'mean': mean.tolist(), 'std': std.tolist(), | |
| 'rows': {'S10': len(data['t0']['uid']), | |
| 'S9': len(data['t0']['uid']) + len(data['t1']['uid']), | |
| 'heldout_full': len(data['heldout']['uid']), 'val': len(selected)}, | |
| 'val_policy': '512 per DE/EN x real/synth cell; half uniform and half score-tail weighted, fixed seed', | |
| 'target_groups': {'emotion': [0, 40], 'voicenet': [40, 97], 'other': [97, 100]}, | |
| } | |
| (OUT / 'manifest.json').write_text(json.dumps(meta, indent=2) + '\n') | |
| print('DONE', meta['rows'], flush=True) | |
| def balanced_val(heldout: dict[str, np.ndarray], n: int) -> np.ndarray: | |
| rng = np.random.default_rng(SEED) | |
| y = heldout['y'] | |
| low = np.nanquantile(y, 0.1, axis=0) | |
| high = np.nanquantile(y, 0.9, axis=0) | |
| extreme = np.isfinite(y) & ((y <= low) | (y >= high)) | |
| # Rarer simultaneous tails get more weight. The same indices serve all four runs. | |
| score = np.sqrt(1.0 + extreme[:, :99].sum(axis=1)) | |
| source_counts = np.bincount(heldout['source'].astype(int)) | |
| score *= 1 / np.sqrt(source_counts[heldout['source'].astype(int)] / source_counts.max()) | |
| out: list[np.ndarray] = [] | |
| assert n % 4 == 0 | |
| for cell in range(4): | |
| candidates = np.flatnonzero(heldout['cell'] == cell) | |
| assert len(candidates) >= n // 4 | |
| ordinary = rng.choice(candidates, n // 8, replace=False) | |
| remaining = np.setdiff1d(candidates, ordinary, assume_unique=True) | |
| p = score[remaining] / score[remaining].sum() | |
| tails = rng.choice(remaining, n // 8, replace=False, p=p) | |
| out.extend([ordinary, tails]) | |
| chosen = np.sort(np.concatenate(out)) | |
| assert len(chosen) == n and len(np.unique(chosen)) == n | |
| return chosen | |
| if __name__ == '__main__': | |
| build() | |