| """Measure codec likelihood and speech quality on fixed reference panels.""" |
| from __future__ import annotations |
| import gc |
| from collections.abc import Sequence |
| from pathlib import Path |
| import numpy as np |
| import soundfile as sf |
| import torch |
| from pydantic import Field |
| from torch import Tensor, nn |
| from torch.nn import functional |
| from transformers import Wav2Vec2CTCTokenizer, Wav2Vec2FeatureExtractor, Wav2Vec2ForCTC, WavLMForXVector |
| from .artifacts import atomic_binary, digest_file, provenance, read_recordings, write_json |
| from .backbone import Generator |
| from .codec import Codec |
| from .configuration import Specification |
| from .dataset import SpeechDataset, collate_layouts |
| from .inference import Sampling, generate_codes |
| from .layout import prompt_layout |
| from .objective import codec_loss |
| from .partition import priority |
| from .waveforms import Waveform, resample_waveform, text_key |
| ASR_ID = 'agbalu/Fadhma-300M' |
| ASR_REVISION = 'd7f5bae148cde28d70d1bcf6df39931c24144583' |
| SPEAKER_ID = 'microsoft/wavlm-base-plus-sv' |
| SPEAKER_REVISION = 'feb593a6c23c1cc3d9510425c29b0a14d2b07b1e' |
| MAX_CER_EXCESS = 0.01 |
| MIN_NATURALNESS = 4.3 |
| MIN_PRONUNCIATION = 4.5 |
| MIN_VOICE_RETENTION = 0.95 |
|
|
| class SpeechObservation(Specification): |
| record_id: str |
| speaker: str |
| reference_characters: int = Field(ge=1) |
| character_errors: int | None = Field(default=None, ge=0) |
| real_character_errors: int | None = Field(default=None, ge=0) |
| voice_cosine: float | None = None |
| real_voice_cosine: float | None = None |
| terminated: bool |
| error: str | None = None |
| waveform_sha256: str | None = None |
|
|
| class ListenerReview(Specification): |
| evaluation_sha256: str |
| naturalness_mean: float = Field(ge=1, le=5) |
| pronunciation_mean: float = Field(ge=1, le=5) |
| voice_retention_fraction: float = Field(ge=0, le=1) |
| expression_retained: bool |
| language_retained: bool |
| ratings: int = Field(ge=2) |
| listeners: int = Field(ge=2) |
|
|
| def edit_count(reference: str, hypothesis: str) -> int: |
| previous = list(range(len(hypothesis) + 1)) |
| for row, left in enumerate(reference, start=1): |
| current = [row] |
| for column, right in enumerate(hypothesis, start=1): |
| current.append(min(current[-1] + 1, previous[column] + 1, previous[column - 1] + (left != right))) |
| previous = current |
| return previous[-1] |
|
|
| def selection_report(observations: Sequence[SpeechObservation], review: ListenerReview | None=None) -> dict[str, object]: |
| support = sum((item.reference_characters for item in observations)) |
| measured = [item for item in observations if item.character_errors is not None and item.real_character_errors is not None] |
| measured_support = sum((item.reference_characters for item in measured)) |
| excess = sum(((item.character_errors or 0) - (item.real_character_errors or 0) for item in measured)) / measured_support if measured_support else None |
| failures = sum((not item.terminated for item in observations)) |
| reasons: list[str] = [] |
| if not observations or failures: |
| reasons.append('empty panel or failed native termination') |
| if len(measured) != len(observations) or excess is None or excess > MAX_CER_EXCESS: |
| reasons.append('CER above matched real-audio floor is missing or exceeds 0.01') |
| voices = [item.voice_cosine for item in observations if item.voice_cosine is not None] |
| if len(voices) != len(observations): |
| reasons.append('speaker verification is incomplete') |
| if review is None: |
| reasons.append('blinded voice/expression/language review is pending') |
| elif review.naturalness_mean < MIN_NATURALNESS or review.pronunciation_mean < MIN_PRONUNCIATION or review.voice_retention_fraction < MIN_VOICE_RETENTION or (not review.expression_retained) or (not review.language_retained): |
| reasons.append('listener quality or retained-capability gate failed') |
| return {'eligible': not reasons, 'reasons': reasons, 'utterances': len(observations), 'reference_characters': support, 'measured_characters': measured_support, 'cer_excess': excess, 'termination_failures': failures, 'mean_voice_cosine': float(np.mean(voices)) if voices else None} |
|
|
| @torch.inference_mode() |
| def heldout_loss(model: Generator, dataset: SpeechDataset, count: int) -> dict[str, object]: |
| device = model.audio_embedding.weight.device |
| was_training = model.training |
| model.train(False) |
| nats = torch.zeros(8, dtype=torch.float64, device=device) |
| counts = torch.zeros(8, dtype=torch.long, device=device) |
| order = sorted(range(len(dataset)), key=lambda index: priority(dataset.recipe.seed, dataset.pairs[index].target))[:count] |
| try: |
| for index in order: |
| text, codes, audio, valid, labels = collate_layouts([dataset.layout(index)], device) |
| with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == 'cuda'): |
| hidden = model.hidden(text, codes, audio, valid) |
| loss = codec_loss(model.classify(hidden), labels) |
| nats += loss.stream_nats.double() |
| counts += loss.stream_counts |
| finally: |
| model.train(was_training) |
| return {'audio_ce_nats': float((nats.sum() / counts.sum()).item()), 'stream_nats': nats.cpu().tolist(), 'stream_counts': counts.cpu().tolist(), 'examples': len(order), 'teacher_forced': True} |
|
|
| def load_control_audio(dataset: SpeechDataset, record_id: str) -> Waveform: |
| entry = dataset.entries[record_id] |
| path = dataset.recipe.output / 'evaluation' / f'{record_id}.wav' |
| if not entry.evaluation_waveform_sha256 or digest_file(path) != entry.evaluation_waveform_sha256: |
| message = f'missing/mismatched real evaluation waveform {path}' |
| raise ValueError(message) |
| values, rate_hz = sf.read(path, dtype='float32') |
| waveform = np.asarray(values, dtype=np.float32) |
| return resample_waveform(waveform, int(rate_hz), 16000) |
|
|
| def extract_input(extractor: Wav2Vec2FeatureExtractor, audio: Waveform, device: torch.device) -> Tensor: |
| features = extractor(audio, sampling_rate=16000, return_tensors='pt') |
| value = features.get('input_values') |
| if not isinstance(value, Tensor): |
| message = 'audio feature extractor returned no input_values tensor' |
| raise TypeError(message) |
| return value.to(device) |
|
|
| @torch.inference_mode() |
| def measure_speech(model: Generator, dataset: SpeechDataset, codec: Codec, output: Path) -> dict[str, object]: |
| output.mkdir(parents=True, exist_ok=True) |
| records = {item.record_id: item for item in read_recordings(dataset.recipe.output / 'recordings.jsonl')} |
| pairs = sorted(dataset.pairs, key=lambda pair: priority(dataset.recipe.seed, pair.target))[:dataset.recipe.evaluation_examples] |
| waves: dict[str, Waveform] = {} |
| failures: dict[str, str] = {} |
| for pair in pairs: |
| target, reference = (dataset.entries[pair.target], dataset.entries[pair.reference]) |
| prompt = prompt_layout(dataset.specials, target.text_ids, dataset.codes(reference), reference.text_ids) |
| try: |
| codes, trace = generate_codes(model, prompt, Sampling(max_new_positions=min(1500, 8192 - prompt.positions))) |
| waveform = codec.decode(codes).numpy() |
| except (ValueError, RuntimeError) as error: |
| failures[pair.target] = str(error) |
| continue |
| path = output / f'{pair.target}.wav' |
| with atomic_binary(path) as stream: |
| sf.write(stream, waveform, 24000, format='WAV', subtype='FLOAT') |
| write_json(path.with_suffix('.stats.json'), {**provenance(dataset.recipe.digest), **trace, 'target': pair.target, 'reference': pair.reference, 'waveform_sha256': digest_file(path)}) |
| waves[pair.target] = resample_waveform(np.asarray(waveform, dtype=np.float32), 24000, 16000) |
| device = model.audio_embedding.weight.device |
| extractor = Wav2Vec2FeatureExtractor.from_pretrained(ASR_ID, revision=ASR_REVISION) |
| tokenizer = Wav2Vec2CTCTokenizer.from_pretrained(ASR_ID, revision=ASR_REVISION) |
| recognizer: nn.Module = Wav2Vec2ForCTC.from_pretrained(ASR_ID, revision=ASR_REVISION) |
| recognizer.to(device).train(False) |
| errors: dict[str, tuple[int, int]] = {} |
| for pair in pairs: |
| if pair.target not in waves: |
| continue |
| hypotheses: list[str] = [] |
| for audio in (waves[pair.target], load_control_audio(dataset, pair.target)): |
| result = recognizer(input_values=extract_input(extractor, audio, device)) |
| logits = getattr(result, 'logits', None) |
| if not isinstance(logits, Tensor): |
| message = 'ASR did not return logits' |
| raise TypeError(message) |
| hypotheses.append(str(tokenizer.decode(logits[0].argmax(-1).cpu().tolist()))) |
| text = text_key(records[pair.target].text).lower() |
| errors[pair.target] = (edit_count(text, text_key(hypotheses[0]).lower()), edit_count(text, text_key(hypotheses[1]).lower())) |
| del recognizer |
| gc.collect() |
| if device.type == 'cuda': |
| torch.cuda.empty_cache() |
| verifier: nn.Module = WavLMForXVector.from_pretrained(SPEAKER_ID, revision=SPEAKER_REVISION) |
| verifier.to(device).train(False) |
| voice_extractor = Wav2Vec2FeatureExtractor.from_pretrained(SPEAKER_ID, revision=SPEAKER_REVISION) |
| observations: list[SpeechObservation] = [] |
| for pair in pairs: |
| text = text_key(records[pair.target].text).lower() |
| if pair.target in failures: |
| observations.append(SpeechObservation(record_id=pair.target, speaker=pair.speaker, reference_characters=len(text), terminated=False, error=failures[pair.target])) |
| continue |
| embeddings: list[Tensor] = [] |
| for audio in (waves[pair.target], load_control_audio(dataset, pair.target), load_control_audio(dataset, pair.reference)): |
| result = verifier(input_values=extract_input(voice_extractor, audio, device)) |
| embedding = getattr(result, 'embeddings', None) |
| if not isinstance(embedding, Tensor): |
| message = 'speaker verifier did not return embeddings' |
| raise TypeError(message) |
| embeddings.append(functional.normalize(embedding.float(), dim=-1)) |
| generated, real, voice_reference = embeddings |
| observations.append(SpeechObservation(record_id=pair.target, speaker=pair.speaker, reference_characters=len(text), character_errors=errors[pair.target][0], real_character_errors=errors[pair.target][1], terminated=True, voice_cosine=float((generated * voice_reference).sum().item()), real_voice_cosine=float((real * voice_reference).sum().item()), waveform_sha256=digest_file(output / f'{pair.target}.wav'))) |
| del verifier |
| report = {**provenance(dataset.recipe.digest), 'observations': [item.model_dump() for item in observations], 'selection': selection_report(observations), 'asr_decoder': 'greedy', 'asr_revision': ASR_REVISION, 'speaker_revision': SPEAKER_REVISION} |
| write_json(output / 'evaluation.stats.json', report) |
| return report |
|
|
