| |
| """Export the actual saved best states and check safe-weight inference on compute.""" |
| import argparse |
| import hashlib |
| import json |
| import os |
| import sys |
| import tempfile |
| from pathlib import Path |
|
|
| import numpy as np |
| import soundfile as sf |
| import torch |
| from safetensors.torch import load_file, save_file |
|
|
| HERE = Path(__file__).resolve().parent |
| GEMINI = Path('/e/scratch/reformo/schuhmann1_moss/whisper_score_regression/gemini_finetune_preparation_20261005') |
| RELEASE = GEMINI.parent / 'humaneness_ears_release_20261007/model' |
|
|
|
|
| def digest(path): |
| h = hashlib.sha256() |
| with Path(path).open('rb') as stream: |
| while part := stream.read(8 * 1024 * 1024): |
| h.update(part) |
| return h.hexdigest() |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--model', choices=('base', 'small'), required=True) |
| args = parser.parse_args() |
| size = args.model |
| release_size = 'medium' if size == 'small' else size |
| torch.set_num_threads(2) |
| source = GEMINI / 'training' / ('whisper_' + size) / 'best.pt' |
| checkpoint = torch.load(source, map_location='cpu', weights_only=False) |
| assert checkpoint['epoch_completed'] == 2 and checkpoint['update'] == 914 |
| norm = json.loads((RELEASE / 'training_normalization.json').read_text()) |
| for key in ['score_names', 'score_mean', 'score_std', 'cps_mean', 'cps_std']: |
| assert checkpoint['normalization'][key] == norm[key], key |
| state = {k: v.detach().contiguous() for k, v in checkpoint['model'].items()} |
| assert all(not k.startswith(('decoder.', 'model.decoder.')) for k in state) |
| output = RELEASE / release_size / 'model.safetensors' |
| save_file(state, output, metadata={'format': 'pt', 'author': 'Christoph Schuhmann', |
| 'license': 'cc-by-4.0', 'checkpoint_sha256': digest(source), |
| 'epochs': '2', 'update': '914', 'teacher': 'gemini-3.8-flash'}) |
| restored = load_file(output) |
| assert set(restored) == set(state) |
| assert all(torch.equal(state[k], restored[k]) for k in state) |
| sys.path.insert(0, str(RELEASE)) |
| from inference import Predictor |
| from model import LayeredMultiTaskWhisper |
| device = 'cuda:' + os.environ.get('LOCAL_RANK', '0') |
| predictor = Predictor(RELEASE, size, device, threads=2) |
| original = LayeredMultiTaskWhisper(RELEASE / release_size, 53, initialize_pretrained=False).to(device).eval() |
| original.load_state_dict(checkpoint['model'], strict=True) |
| assert sum(p.numel() for p in original.parameters()) == json.loads( |
| (RELEASE / 'training/configs' / (size + '_recorded_run.json')).read_text())['trainable_total_parameters'] |
| |
| generator = torch.Generator().manual_seed(20261007) |
| mel = torch.randn((1, 80, 100), generator=generator).to(device) |
| mask = torch.ones((1, 100), dtype=torch.long, device=device) |
| with torch.inference_mode(): |
| a = original(mel, mask, predict_events=True, max_pred_events=32) |
| b = predictor.model(mel, mask, predict_events=True, max_pred_events=32) |
| assert set(a) == set(b) |
| assert all(torch.equal(a[k], b[k]) for k in a), 'Export altered inference outputs' |
| del original, checkpoint, state, restored, a, b |
| |
| row = None |
| with (GEMINI / 'prepared/targets.jsonl').open() as stream: |
| for line in stream: |
| candidate = json.loads(line) |
| if candidate['split'] == 'test' and candidate['ready_for_whisper_training']: |
| row = candidate |
| break |
| sys.path.insert(0, str(RELEASE / 'training/gemini_finetune')) |
| from common import TarReader |
| wave = TarReader().audio(row, verify=True) |
| with tempfile.TemporaryDirectory(prefix='whisper-release-') as directory: |
| path = Path(directory) / 'heldout.wav' |
| sf.write(path, wave, 16000, subtype='FLOAT') |
| result = predictor.predict(path, include_frames=True) |
| assert len(result['scores_raw']) == len(result['scores_z']) == len(result['scores_training_z']) == 192 |
| assert len(result['orange_timbre_128d']) == 128 and len(result['orange_identity_250d']) == 250 |
| assert len(result['burst_probability_per_20ms_frame']) > 0 |
| json.dumps(result, allow_nan=False) |
| for event in result['binary_burst_regions'] + result['burst_event_proposals']: |
| assert 0 <= event['start_s'] < event['end_s'] <= result['duration_s'] + 1e-6 |
| for event in result['burst_event_proposals']: |
| assert len(event['top3_classes']) == 3 |
| |
| cpu = Predictor(RELEASE, size, 'cpu', threads=1) |
| short = Path(directory) / 'heldout-short.wav' |
| sf.write(short, wave[:16000], 16000, subtype='FLOAT') |
| cpu_result = cpu.predict(short) |
| json.dumps(cpu_result, allow_nan=False) |
| assert len(cpu_result['scores_raw']) == 192 |
| evidence = {'model': size, 'checkpoint_epoch': 2, 'updates': 914, |
| 'source_checkpoint_sha256': digest(source), 'safe_weights_sha256': digest(output), |
| 'safe_weights_bytes': output.stat().st_size, |
| 'all_tensors_bitwise_equal': True, 'all_model_outputs_bitwise_equal': True, |
| 'heldout_audio_sha256': row['sha256_mp3'], 'heldout_duration_s': result['duration_s'], |
| 'targets': 192, 'speaker_dimensions': [128, 250], 'finite_outputs': True, |
| 'timestamps_have_start_and_end': True, 'device': device, |
| 'one_thread_cpu_inference_passed': True, |
| 'torch': torch.__version__, 'slurm_job': os.environ.get('SLURM_JOB_ID')} |
| path = RELEASE / 'verification' / (size + '.json') |
| path.parent.mkdir(exist_ok=True) |
| path.write_text(json.dumps(evidence, indent=2) + '\n') |
| print(json.dumps(evidence), flush=True) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
