| """Audit embedded source audio through resumable shard transactions.""" |
| from __future__ import annotations |
| import hashlib |
| import io |
| import json |
| import threading |
| from collections.abc import Iterator |
| from datetime import date |
| from pathlib import Path |
| from typing import BinaryIO |
| import numpy as np |
| import pyarrow.parquet as pq |
| import soundfile as sf |
| from pydantic import Field |
| from .artifacts import Recording, Rejection, atomic_binary, digest_file, provenance, read_recordings, run_lock, verify_conservation, write_json |
| from .configuration import Recipe, Specification |
| from .pairs import pair_recordings |
| from .partition import partition_recordings |
| from .waveforms import Waveform, measure_signal, resample_waveform, text_key |
| from .progress import Tracker, stage |
| DECODE_LOCK = threading.Lock() |
|
|
| class SourceFile(Specification): |
| path: str |
| sha256: str |
|
|
| class Acquisition(Specification): |
| revision: str |
| files: tuple[SourceFile, ...] |
| access_date: date |
|
|
| class ShardResult(Specification): |
| config_hash: str |
| source_sha256: str |
| rows: int = Field(ge=0) |
| accepted: int = Field(ge=0) |
| rejected: int = Field(ge=0) |
| recordings_sha256: str |
| rejects_sha256: str |
| provenance: dict[str, object] |
|
|
| def read_acquisition(recipe: Recipe) -> Acquisition: |
| path = recipe.source / 'acquisition.stats.json' |
| payload: object = json.loads(path.read_bytes()) |
| if not isinstance(payload, dict) or not isinstance(payload.get('files'), list): |
| message = f'invalid acquisition file list: {path}' |
| raise TypeError(message) |
| files: list[SourceFile] = [] |
| for item in payload['files']: |
| if not isinstance(item, dict) or not isinstance(item.get('path'), str): |
| message = f'invalid acquisition entry at {path}: {item!r}' |
| raise TypeError(message) |
| if item['path'].endswith('.parquet'): |
| files.append(SourceFile.model_validate({'path': item['path'], 'sha256': item.get('sha256')})) |
| acquisition = Acquisition.model_validate({'revision': payload.get('revision'), 'files': tuple(files), 'access_date': payload.get('access_date')}) |
| if acquisition.revision != recipe.source_revision or not files: |
| message = f'source identity mismatch or no Parquet files at {path}: {acquisition.revision}' |
| raise ValueError(message) |
| return acquisition |
|
|
| def source_files(recipe: Recipe) -> tuple[Path, ...]: |
| path = recipe.source / 'acquisition.stats.json' |
| acquisition = read_acquisition(recipe) |
| expected: dict[str, str] = {} |
| for item in acquisition.files: |
| name = Path(item.path).name |
| if name in expected: |
| message = f'duplicate acquisition shard name {name!r} in {path}' |
| raise ValueError(message) |
| expected[name] = item.sha256 |
| actual = tuple(sorted((recipe.source / 'data').glob('*.parquet'))) |
| if set(expected) != {item.name for item in actual}: |
| message = f'source shard set differs from {path}: expected={sorted(expected)} actual={actual}' |
| raise ValueError(message) |
| for shard in actual: |
| observed = digest_file(shard) |
| if observed != expected[shard.name]: |
| message = f'checksum mismatch {shard}: observed={observed}, expected={expected[shard.name]}' |
| raise ValueError(message) |
| return actual |
|
|
| def raw_rows(shard: Path) -> Iterator[tuple[int, object]]: |
| with pq.ParquetFile(shard) as parquet: |
| offset = 0 |
| for batch in parquet.iter_batches(batch_size=128, columns=['audio', 'transcription', 'client_id'], use_threads=False): |
| for value in batch.to_pylist(): |
| yield (offset, value) |
| offset += 1 |
| verify_conservation(parquet.metadata.num_rows, offset, 0) |
|
|
| def embedded_bytes(value: object, max_size_bytes: int) -> tuple[bytes, str, str]: |
| if not isinstance(value, dict): |
| message = f'expected audio record, got {type(value).__name__}' |
| raise TypeError(message) |
| audio, text, speaker = (value.get('audio'), value.get('transcription'), value.get('client_id')) |
| if not isinstance(audio, dict) or not isinstance(audio.get('bytes'), bytes): |
| message = 'missing embedded audio.bytes' |
| raise TypeError(message) |
| raw = audio['bytes'] |
| if len(raw) > max_size_bytes or not raw: |
| message = f'embedded audio size_bytes={len(raw)}, limit={max_size_bytes}' |
| raise ValueError(message) |
| if not isinstance(text, str) or not text_key(text): |
| message = f'missing/empty transcription: {text!r}' |
| raise ValueError(message) |
| if not isinstance(speaker, str) or not speaker.strip(): |
| message = f'missing/empty client_id: {speaker!r}' |
| raise ValueError(message) |
| return (raw, text, speaker) |
|
|
| def decode_audio(raw: bytes) -> tuple[Waveform, int]: |
| with DECODE_LOCK: |
| waveform, rate_hz = sf.read(io.BytesIO(raw), dtype='float32', always_2d=True) |
| if waveform.shape[1] != 1: |
| message = f'expected mono audio, got channels={waveform.shape[1]}' |
| raise ValueError(message) |
| mono = np.asarray(waveform[:, 0], dtype=np.float32) |
| measure_signal(mono, rate_hz) |
| return (mono, int(rate_hz)) |
|
|
| def recording_for(value: object, shard: Path, row: int, recipe: Recipe, retrieval_date: str) -> Recording: |
| raw, text, speaker = embedded_bytes(value, recipe.max_embedded_bytes) |
| waveform, rate_hz = decode_audio(raw) |
| signal = measure_signal(waveform, rate_hz) |
| duration_s = signal.samples / rate_hz |
| if not recipe.min_duration_s <= duration_s <= recipe.max_duration_s: |
| message = f'duration_s={duration_s} outside [{recipe.min_duration_s},{recipe.max_duration_s}]' |
| raise ValueError(message) |
| converted = resample_waveform(waveform, rate_hz, recipe.rate_hz) |
| measure_signal(converted, recipe.rate_hz) |
| identity = f'{recipe.source_revision}:{shard.name}:{row}' |
| return Recording(record_id=hashlib.sha256(identity.encode()).hexdigest(), shard=shard.name, row=row, audio_sha256=hashlib.sha256(raw).hexdigest(), speaker=speaker, text=text, text_key=text_key(text), stored_rate_hz=rate_hz, samples=signal.samples, prepared_samples=converted.size, duration_s=duration_s, peak=signal.peak, rms_dbfs=signal.rms_dbfs, clipping_fraction=signal.clipping_fraction, source_id=recipe.source_id, source_revision=recipe.source_revision, license=recipe.license, retrieval_date=retrieval_date) |
|
|
| def write_line(stream: BinaryIO, value: Specification) -> None: |
| stream.write((value.model_dump_json() + '\n').encode()) |
|
|
| def prepare_shard(shard: Path, recipe: Recipe, date: str) -> ShardResult: |
| root = recipe.output / 'shards' / shard.stem |
| stats_path = root / 'audit.stats.json' |
| accepted_path, rejects_path = (root / 'recordings.jsonl', root / 'rejects.jsonl') |
| source_digest = digest_file(shard) |
| if stats_path.is_file(): |
| result = ShardResult.model_validate_json(stats_path.read_bytes()) |
| if result.config_hash != recipe.digest or result.source_sha256 != source_digest or result.recordings_sha256 != digest_file(accepted_path) or (result.rejects_sha256 != digest_file(rejects_path)): |
| message = f'prepared artifact identity/checksum mismatch: {stats_path}' |
| raise ValueError(message) |
| verify_conservation(result.rows, result.accepted, result.rejected) |
| return result |
| rows = accepted = rejected = 0 |
| with atomic_binary(accepted_path) as kept, atomic_binary(rejects_path) as rejects: |
| for row, value in raw_rows(shard): |
| rows += 1 |
| try: |
| record = recording_for(value, shard, row, recipe, date) |
| except (ValueError, TypeError, sf.LibsndfileError) as error: |
| rejected += 1 |
| rejection = Rejection(record_id=f'{shard.name}:{row}', shard=shard.name, row=row, reason=str(error)) |
| write_line(rejects, rejection) |
| else: |
| accepted += 1 |
| write_line(kept, record) |
| verify_conservation(rows, accepted, rejected) |
| result = ShardResult(config_hash=recipe.digest, source_sha256=source_digest, rows=rows, accepted=accepted, rejected=rejected, recordings_sha256=digest_file(accepted_path), rejects_sha256=digest_file(rejects_path), provenance=provenance(recipe.digest)) |
| write_json(stats_path, result.model_dump()) |
| return result |
|
|
| def prepare_audio(recipe: Recipe) -> dict[str, object]: |
| with run_lock(recipe.output), stage('decode source audio'): |
| shards = source_files(recipe) |
| tracker = Tracker('decoded shards', len(shards), interval_s=5) |
| counts: list[ShardResult] = [] |
| retrieval_date = read_acquisition(recipe).access_date.isoformat() |
| with atomic_binary(recipe.output / 'recordings.jsonl') as combined: |
| for shard in shards: |
| counts.append(prepare_shard(shard, recipe, retrieval_date)) |
| path = recipe.output / 'shards' / shard.stem / 'recordings.jsonl' |
| with path.open('rb') as stream: |
| while (block := stream.read(1 << 20)): |
| combined.write(block) |
| tracker.advance() |
| report: dict[str, object] = {**provenance(recipe.digest), 'input_rows': sum((item.rows for item in counts)), 'accepted_rows': sum((item.accepted for item in counts)), 'rejected_rows': sum((item.rejected for item in counts)), 'language_reviewed': False, 'recordings_sha256': digest_file(recipe.output / 'recordings.jsonl'), 'shards': [item.model_dump() for item in counts]} |
| write_json(recipe.output / 'preparation.stats.json', report) |
| return report |
|
|
| def prepare_pairs(recipe: Recipe) -> dict[str, object]: |
| with run_lock(recipe.output), stage('partition reference pairs'): |
| records = tuple(read_recordings(recipe.output / 'recordings.jsonl')) |
| partition = partition_recordings(records, recipe) |
| pairing = pair_recordings(records, partition, recipe.seed) |
| for split in ('train', 'validation', 'test'): |
| if not any((pair.split == split for pair in pairing.pairs)): |
| message = f'no eligible pairs in {split}: inspect speaker/script coverage' |
| raise ValueError(message) |
| with atomic_binary(recipe.output / 'pairs.jsonl') as stream: |
| for pair in pairing.pairs: |
| write_line(stream, pair) |
| excluded = {**partition.excluded, **pairing.excluded} |
| write_json(recipe.output / 'pair-rejects.json', excluded) |
| verify_conservation(len(records), len(pairing.pairs), len(excluded)) |
| report: dict[str, object] = {**provenance(recipe.digest), 'input_rows': len(records), 'paired_rows': len(pairing.pairs), 'excluded_rows': len(excluded), 'split_pairs': {split: sum((pair.split == split for pair in pairing.pairs)) for split in ('train', 'validation', 'test')}, 'recordings_sha256': digest_file(recipe.output / 'recordings.jsonl'), 'pairs_sha256': digest_file(recipe.output / 'pairs.jsonl')} |
| write_json(recipe.output / 'pairs.stats.json', report) |
| return report |
|
|
