| """Cache immutable codec arrays for deterministic reference-conditioned batches.""" |
| from __future__ import annotations |
| import json |
| from collections.abc import Iterator, Mapping, Sequence |
| from pathlib import Path |
| from typing import Protocol |
| import numpy as np |
| import soundfile as sf |
| import torch |
| from pydantic import Field |
| from torch import Tensor |
| from torch.utils.data import Dataset |
| from .artifacts import Pair, Split, atomic_binary, digest_file, provenance, read_recordings, run_lock, verify_conservation, write_json |
| from .audio import decode_audio, embedded_bytes, raw_rows, source_files |
| from .codec import Codec |
| from .configuration import Recipe, Specification |
| from .layout import CODEBOOKS, SPECIALS, SequenceLayout, prompt_layout, target_layout |
| from .partition import priority |
| from .waveforms import resample_waveform |
| from .progress import Tracker |
|
|
| class TextTokenizer(Protocol): |
|
|
| def encode(self, text: str, *, add_special_tokens: bool=False) -> list[int]: |
| ... |
|
|
| def get_added_vocab(self) -> dict[str, int]: |
| ... |
|
|
| class TokenEntry(Specification): |
| record_id: str |
| file: str |
| offset_frames: int = Field(ge=0) |
| frames: int = Field(ge=1) |
| text_ids: tuple[int, ...] = Field(min_length=1) |
| audio_sha256: str |
| duration_s: float = Field(gt=0) |
| padding_samples: int = Field(ge=0) |
| evaluation_waveform_sha256: str | None = None |
|
|
| def cache_codes(recipe: Recipe, codec: Codec, tokenizer: TextTokenizer) -> dict[str, object]: |
| with run_lock(recipe.output): |
| records = {record.record_id: record for record in read_recordings(recipe.output / 'recordings.jsonl')} |
| by_location = {(record.shard, record.row): record for record in records.values()} |
| with (recipe.output / 'pairs.jsonl').open() as pairs_stream: |
| pairs = tuple((Pair.model_validate_json(line) for line in pairs_stream)) |
| protected: set[str] = set() |
| for split in ('validation', 'test'): |
| candidates = sorted((pair for pair in pairs if pair.split == split), key=lambda pair: priority(recipe.seed, pair.target)) |
| for pair in candidates[:recipe.evaluation_examples]: |
| protected.update((pair.target, pair.reference)) |
| shards = source_files(recipe) |
| tracker = Tracker('codec shards', len(shards), interval_s=5) |
| files: list[dict[str, object]] = [] |
| total = 0 |
| with atomic_binary(recipe.output / 'tokens.jsonl') as index: |
| for shard in shards: |
| token_path = recipe.output / 'tokens' / f'{shard.stem}.codes' |
| entry_path = token_path.with_suffix('.jsonl') |
| stats_path = token_path.with_suffix('.stats.json') |
| identity = {'config_hash': recipe.digest, 'source_sha256': digest_file(shard), 'recordings_sha256': digest_file(recipe.output / 'recordings.jsonl')} |
| if stats_path.is_file(): |
| report = json.loads(stats_path.read_bytes()) |
| if any((report.get(key) != value for key, value in identity.items())) or report.get('token_sha256') != digest_file(token_path) or report.get('index_sha256') != digest_file(entry_path): |
| message = f'codec cache identity/checksum mismatch at {stats_path}' |
| raise ValueError(message) |
| else: |
| offset = count = 0 |
| with atomic_binary(token_path) as binary, atomic_binary(entry_path) as entries: |
| for row, value in raw_rows(shard): |
| recording = by_location.get((shard.name, row)) |
| if recording is None: |
| continue |
| raw, text, _ = embedded_bytes(value, recipe.max_embedded_bytes) |
| waveform, rate_hz = decode_audio(raw) |
| converted = resample_waveform(waveform, rate_hz) |
| waveform_digest: str | None = None |
| if recording.record_id in protected: |
| waveform_path = recipe.output / 'evaluation' / f'{recording.record_id}.wav' |
| with atomic_binary(waveform_path) as waveform_stream: |
| sf.write(waveform_stream, converted, 24000, format='WAV', subtype='FLOAT') |
| waveform_digest = digest_file(waveform_path) |
| codes = codec.encode(converted) |
| values = np.asarray(codes.numpy(), dtype='<u2') |
| binary.write(values.tobytes()) |
| entry = TokenEntry(record_id=recording.record_id, file=token_path.name, offset_frames=offset, frames=codes.shape[0], text_ids=tuple(tokenizer.encode(text, add_special_tokens=False)), audio_sha256=recording.audio_sha256, duration_s=recording.duration_s, padding_samples=max(0, 24000 - converted.size), evaluation_waveform_sha256=waveform_digest) |
| entries.write((entry.model_dump_json() + '\n').encode()) |
| offset += codes.shape[0] |
| count += 1 |
| report = {**provenance(recipe.digest), **identity, 'records': count, 'frames': offset, 'token_sha256': digest_file(token_path), 'index_sha256': digest_file(entry_path)} |
| write_json(stats_path, report) |
| with entry_path.open('rb') as stream: |
| while (block := stream.read(1 << 20)): |
| index.write(block) |
| total += int(report['records']) |
| files.append({'file': token_path.name, 'sha256': report['token_sha256']}) |
| tracker.advance() |
| verify_conservation(len(records), total, 0) |
| report = {**provenance(recipe.digest), 'records': total, 'files': files, 'index_sha256': digest_file(recipe.output / 'tokens.jsonl')} |
| write_json(recipe.output / 'tokens.stats.json', report) |
| return report |
|
|
| class SpeechDataset(Dataset[SequenceLayout]): |
|
|
| def __init__(self, recipe: Recipe, specials: Mapping[str, int], split: Split) -> None: |
| self.recipe = recipe |
| self.specials = dict(specials) |
| if set(SPECIALS) - self.specials.keys(): |
| message = 'tokenizer is missing native TTS specials' |
| raise ValueError(message) |
| report = json.loads((recipe.output / 'tokens.stats.json').read_bytes()) |
| if report.get('config_hash') != recipe.digest or report.get('index_sha256') != digest_file(recipe.output / 'tokens.jsonl'): |
| message = 'codec index identity/checksum mismatch' |
| raise ValueError(message) |
| for item in report['files']: |
| path = recipe.output / 'tokens' / item['file'] |
| if Path(item['file']).name != item['file'] or digest_file(path) != item['sha256']: |
| message = f'codec binary identity/checksum mismatch: {path}' |
| raise ValueError(message) |
| with (recipe.output / 'tokens.jsonl').open() as stream: |
| self.entries = {entry.record_id: entry for line in stream if (entry := TokenEntry.model_validate_json(line))} |
| with (recipe.output / 'pairs.jsonl').open() as stream: |
| candidates = tuple((pair for line in stream if (pair := Pair.model_validate_json(line)).split == split)) |
| accepted: list[Pair] = [] |
| excluded: list[dict[str, object]] = [] |
| for pair in candidates: |
| target, reference = (self.entries[pair.target], self.entries[pair.reference]) |
| positions = mixed_positions(target, reference) |
| if positions > recipe.max_mixed_positions: |
| excluded.append({'target': pair.target, 'reference': pair.reference, 'positions': positions, 'reason': 'mixed context exceeds maximum'}) |
| else: |
| accepted.append(pair) |
| self.pairs = tuple(accepted) |
| verify_conservation(len(candidates), len(self.pairs), len(excluded)) |
| write_json(recipe.output / f'{split}-context-rejects.json', excluded) |
| write_json(recipe.output / f'{split}-context.stats.json', {**provenance(recipe.digest), 'input_pairs': len(candidates), 'accepted_pairs': len(self.pairs), 'rejected_pairs': len(excluded)}) |
| if not self.pairs: |
| message = f'empty codec dataset split={split}' |
| raise ValueError(message) |
| self.identity = digest_file(recipe.output / 'tokens.stats.json') + digest_file(recipe.output / 'pairs.jsonl') |
|
|
| def __len__(self) -> int: |
| return len(self.pairs) |
|
|
| def __getitem__(self, index: int) -> SequenceLayout: |
| return self.layout(index) |
|
|
| def positions(self, index: int) -> int: |
| pair = self.pairs[index] |
| return mixed_positions(self.entries[pair.target], self.entries[pair.reference]) |
|
|
| def codes(self, entry: TokenEntry) -> Tensor: |
| path = self.recipe.output / 'tokens' / entry.file |
| if Path(entry.file).name != entry.file: |
| message = f'invalid codec file name={entry.file!r} for {entry.record_id}' |
| raise ValueError(message) |
| byte_count = entry.frames * CODEBOOKS * 2 |
| with path.open('rb') as stream: |
| stream.seek(entry.offset_frames * CODEBOOKS * 2) |
| raw = stream.read(byte_count) |
| if len(raw) != byte_count: |
| message = f'short codec read {path}:{entry.offset_frames}: {len(raw)}/{byte_count} bytes' |
| raise ValueError(message) |
| values = np.frombuffer(raw, dtype='<u2').astype(np.int64).reshape(entry.frames, CODEBOOKS) |
| return torch.from_numpy(values) |
|
|
| def layout(self, index: int) -> SequenceLayout: |
| pair = self.pairs[index] |
| target, reference = (self.entries[pair.target], self.entries[pair.reference]) |
| prompt = prompt_layout(self.specials, target.text_ids, self.codes(reference), reference.text_ids) |
| return target_layout(prompt, self.codes(target), self.recipe.max_mixed_positions) |
|
|
| def windows(self) -> Iterator[tuple[int, ...]]: |
| order = sorted(range(len(self.pairs)), key=lambda index: priority(self.recipe.seed, self.pairs[index].target)) |
| for start in range(0, len(order), self.recipe.update_examples): |
| window = order[start:start + self.recipe.update_examples] |
| yield tuple(sorted(window, key=self.positions)) |
|
|
| def collate_layouts(layouts: Sequence[SequenceLayout], device: torch.device) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: |
| if not layouts: |
| message = 'cannot collate an empty codec microbatch' |
| raise ValueError(message) |
| longest = max((layout.positions for layout in layouts)) |
| text = torch.zeros((len(layouts), longest), dtype=torch.long) |
| codes = torch.zeros((len(layouts), longest, CODEBOOKS), dtype=torch.long) |
| audio = torch.zeros_like(text, dtype=torch.bool) |
| valid = torch.zeros_like(text, dtype=torch.bool) |
| labels = torch.full_like(codes, -100) |
| for row, layout in enumerate(layouts): |
| size = layout.positions |
| text[row, :size], codes[row, :size] = (layout.text_ids, layout.audio_codes) |
| audio[row, :size], labels[row, :size] = (layout.audio_mask, layout.labels) |
| valid[row, :size] = True |
| values: tuple[Tensor, ...] = (text, codes, audio, valid, labels) |
| if device.type == 'cuda': |
| values = tuple((value.pin_memory() for value in values)) |
| moved = tuple((value.to(device, non_blocking=device.type == 'cuda') for value in values)) |
| return (moved[0], moved[1], moved[2], moved[3], moved[4]) |
|
|
| def collate_cpu(layouts: Sequence[SequenceLayout]) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: |
| return collate_layouts(layouts, torch.device('cpu')) |
|
|
| def mixed_positions(target: TokenEntry, reference: TokenEntry) -> int: |
| return 5 + len(target.text_ids) + len(reference.text_ids) + target.frames + reference.frames + 2 * (CODEBOOKS - 1) |
|
|
