#!/usr/bin/env python3 """Evaluate Gemini-tuned Whisper and the original weights on the same Flash holdout.""" import argparse import json import os import sys from concurrent.futures import ThreadPoolExecutor import numpy as np import torch from torch.utils.data import DataLoader from transformers import WhisperFeatureExtractor from study_paths import ROOT, CODE, RELEASE, GEMINI, BENCH, read, write from predict_benchmarks import score_metrics from matched_acted import run as acted from train_probes import evaluate sys.path.insert(0, str(CODE / 'gemini_finetune')) from data import GeminiDataset, GeminiCollator sys.path.insert(0, str(CODE)) from layered_multitask_model import LayeredMultiTaskWhisper from evaluate_public_benchmarks import decode, parquet_samples, emolia_samples class P3WhisperData: """Exactly the same deterministic 2k P3 split selections as the probes.""" def __init__(self,split): from study_data import SourceData self.source=SourceData('p3_'+split) def __len__(self):return len(self.source) def __getitem__(self,index): key,wave,meta,targets=self.source[index] return {**targets,'key':key,'uid':key,'domain':'gemini','enabled':True,'audio':wave,'cps':None, 'timbre_valid':bool(targets['timbre_valid']),'identity_valid':bool(targets['identity_valid']), 'events_frames':list(zip(targets['event_starts'],targets['event_ends'])), 'event_class_ids':targets['event_classes']} class HoldoutLoader: def __init__(self, model_path, norm): self.collator = GeminiCollator(model_path, norm) def __call__(self, data): return DataLoader(data, batch_size=32, num_workers=4, multiprocessing_context='spawn', collate_fn=self.collator) def forward(model, batch): with torch.autocast('cuda', torch.bfloat16): oracle = model(batch['mel'], batch['mel_mask'], batch['event_starts'], batch['event_ends']) predicted = model(batch['mel'], batch['mel_mask'], predict_events=True, max_pred_events=32) for key in ('predicted_event_starts', 'predicted_event_ends', 'predicted_event_valid'): oracle[key] = predicted[key] oracle['predicted_event_class'] = predicted['event_class'] return oracle def run(size): device = torch.device('cuda', int(os.environ.get('LOCAL_RANK', 0))) torch.cuda.set_device(device) torch.set_num_threads(2) classes = read(RELEASE / 'classes.json')['names'] cfg = read(CODE / 'gemini_finetune/configs' / ('whisper_' + size + '.json')) checkpoint = GEMINI / 'training' / ('whisper_' + size) / 'best.pt' out = ROOT / 'whisper' / 'gemini' / size out.mkdir(parents=True, exist_ok=True) if (out / 'COMPLETE.json').exists(): return norm = read(RELEASE / 'training_normalization.json') model = LayeredMultiTaskWhisper(RELEASE / size, len(classes), initialize_pretrained=False).to(device) holdout = GeminiDataset(GEMINI / 'prepared', 'test') datasets={'flash':holdout,'p3_validation':P3WhisperData('validation'),'p3_test':P3WhisperData('test')} for phase,path in [('legacy',cfg['checkpoint']),('gemini',checkpoint)]: model.load_state_dict(torch.load(path,map_location='cpu',weights_only=False)['model'],strict=True) for domain,data in datasets.items(): filename=('legacy_on_flash_test_metrics.json' if phase=='legacy' else 'gemini_test_metrics.json') if domain=='flash' else phase+'_'+domain+'_metrics.json' target=out/filename if not target.exists(): metrics=evaluate(model,data,device,torch.ones(len(classes),device=device),True, HoldoutLoader(RELEASE/size,norm),forward) metrics.update(domain=domain,phase=phase,split_scope='Fixed identical held-out audio selection before and after fine-tuning') write(target,metrics) model.load_state_dict(torch.load(checkpoint, map_location='cpu', weights_only=False)['model'], strict=True) model.eval() extractor = WhisperFeatureExtractor.from_pretrained(RELEASE / size, local_files_only=True) predictions = {} for kind in read(ROOT / 'study.json')['benchmarks']: stem = out / (kind + '_predictions') if stem.with_suffix('.npz').exists() and stem.with_suffix('.jsonl').exists(): scores = np.load(stem.with_suffix('.npz'))['scores'] rows = [json.loads(line) for line in stem.with_suffix('.jsonl').open()] else: samples = parquet_samples(kind) if kind in ('emonet', 'crema', 'ravdess') else emolia_samples(kind) rows, scores = [], [] with ThreadPoolExecutor(4) as pool, torch.inference_mode(): for offset in range(0, len(samples), 32): selected = samples[offset:offset + 32] audio = list(pool.map(decode, selected)) features = extractor([w for w,_,_ in audio], sampling_rate=16000, padding='longest', truncation=True, max_length=480000, return_attention_mask=True, return_tensors='np') mel, mask = features['input_features'], features['attention_mask'] if mel.shape[-1] % 2: mel = np.pad(mel, ((0,0),(0,0),(0,1))) mask = np.pad(mask, ((0,0),(0,1))) with torch.autocast('cuda', torch.bfloat16): result = model(torch.from_numpy(mel).to(device), torch.from_numpy(mask).to(device)) scores.append(result['scores'].float().cpu().numpy() * np.asarray(norm['score_std']) + np.asarray(norm['score_mean'])) rows.extend({**s.metadata, 'key': s.key, 'duration_s': d, 'truncated_30s': t} for s,(_,d,t) in zip(selected,audio)) scores = np.concatenate(scores).astype(np.float32) np.savez_compressed(stem.with_suffix('.npz'), scores=scores) stem.with_suffix('.jsonl').write_text(''.join(json.dumps(row)+'\n' for row in rows)) predictions[kind] = scores, rows if kind in ('crema', 'ravdess'): acted(kind, scores, rows, out) score_metrics(predictions, 'whisper_' + size, out) write(out / 'COMPLETE.json', {'checkpoint': str(checkpoint), 'test_clips': len(holdout), 'p3_validation_clips':len(datasets['p3_validation']),'p3_test_clips':len(datasets['p3_test']), 'benchmarks': list(predictions)}) if __name__ == '__main__': ap = argparse.ArgumentParser(); ap.add_argument('--model', choices=['base','small','all'], required=True) args=ap.parse_args() if args.model!='all': run(args.model) else: rank=int(os.environ.get('LOCAL_RANK',0)) if rank<2: run(('base','small')[rank]) else: size=('base','small')[rank-2];stage='S4' if size=='base' else 'S3' out=ROOT/'whisper/legacy'/size;out.mkdir(parents=True,exist_ok=True) for kind in ('crema','ravdess'): stem=BENCH/'predictions'/(kind+'_'+size+'_'+stage) acted(kind,np.load(stem.with_suffix('.npz'))['scores'], [json.loads(line) for line in stem.with_suffix('.jsonl').open()],out)