| """Latent dataset for training: on-disk format, loader and rejection report. |
| |
| On-disk format (a *latent dataset directory*):: |
| |
| <dir>/index.jsonl one JSON object per segment (UTF-8, one per line) |
| <dir>/shard-00000.npy float16 or float32 array [total_frames, latent_dim] |
| <dir>/shard-00001.npy ... |
| |
| Each ``index.jsonl`` record has exactly these keys: |
| |
| ================== ========================================================================== |
| ``segment_id`` non-empty string, unique across the index |
| ``shard`` file name of a shard in the same directory (no path separators) |
| ``offset`` int >= 0, first frame row of the segment inside the shard |
| ``num_frames`` int >= 1, number of latent frames; ``offset + num_frames <= rows`` |
| ``text`` tag-spec tagged Turkish text (parsed by :func:`nida.model.vocab.encode_text`) |
| ``emotion`` one of :data:`EMOTIONS` or ``null`` |
| ``intensity`` float in [0, 1] or ``null`` |
| ``rate_bucket`` int in 0..4 or ``null`` |
| ``loudness_bucket`` int in 0..2 or ``null`` |
| ``group`` non-empty string (speaker / source video id; splits never cross groups) |
| ``split`` ``train`` | ``val`` | ``test`` |
| ================== ========================================================================== |
| |
| Optional key: ``language`` — one of :data:`LANGUAGES` (ADR-0007); ``"tr"`` when absent, so |
| datasets written before it stay valid. |
| |
| ``null`` condition values map to the model's *null* embedding (the id one past the last real |
| class). A segment of ``n`` frames is the latent slice ``shard[offset : offset + n]``; all |
| shards share one ``latent_dim``. Shards are opened with ``np.load(mmap_mode="r")`` and |
| converted to float32 only for the rows of a requested segment. |
| |
| Structural problems (bad JSON, missing/extra keys, wrong types, duplicate ids, missing |
| shard, out-of-range rows, inconsistent ``latent_dim``, a ``group`` in more than one |
| split) raise :class:`DataFormatError` -- |
| they mean the producer is broken. *Item-level* problems that are properties of the data are |
| **rejected with a counted reason** and reported in :class:`LoadReport` (never dropped |
| silently): ``text_longer_than_frames`` (the F5-style text padding needs |
| ``len(text_ids) <= num_frames``), ``text_encoding_error`` (unknown character / bad tag) and |
| ``empty_text``. |
| |
| Several directories (one per data-pilot root) can be trained on as one union with |
| :func:`load_latent_datasets`; see its docstring for the cross-directory checks. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| import re |
| from collections import Counter |
| from collections.abc import Mapping, Sequence |
| from dataclasses import dataclass, field, replace |
| from pathlib import Path |
|
|
| import numpy as np |
| import numpy.typing as npt |
|
|
| from nida.data.timings import SegmentTimings, load_timings |
| from nida.model.vocab import encode_text |
|
|
| __all__ = [ |
| "EMOTIONS", |
| "LANGUAGES", |
| "N_LOUDNESS_BUCKETS", |
| "N_RATE_BUCKETS", |
| "REJECT_EMPTY_TEXT", |
| "REJECT_TEXT_ENCODING", |
| "REJECT_TEXT_TOO_LONG", |
| "SPLITS", |
| "DataFormatError", |
| "LatentDataset", |
| "LatentStats", |
| "LoadReport", |
| "Rejection", |
| "Segment", |
| "compute_latent_stats", |
| "load_latent_dataset", |
| "load_latent_datasets", |
| ] |
|
|
| |
| |
| LANGUAGES: tuple[str, ...] = ("tr", "en") |
|
|
| |
| |
| EMOTIONS: tuple[str, ...] = ( |
| "angry", |
| "disgusted", |
| "fearful", |
| "happy", |
| "neutral", |
| "sad", |
| "surprised", |
| ) |
| N_RATE_BUCKETS = 5 |
| N_LOUDNESS_BUCKETS = 3 |
| SPLITS: tuple[str, ...] = ("train", "val", "test") |
|
|
| REJECT_TEXT_TOO_LONG = "text_longer_than_frames" |
| REJECT_TEXT_ENCODING = "text_encoding_error" |
| REJECT_EMPTY_TEXT = "empty_text" |
| REJECT_NO_TIMINGS = "no_timings" |
| REJECT_TIMINGS_MISMATCH = "timings_mismatch" |
|
|
| INDEX_NAME = "index.jsonl" |
| _SHARD_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*\.npy$") |
| _RECORD_KEYS = ( |
| "segment_id", |
| "shard", |
| "offset", |
| "num_frames", |
| "text", |
| "emotion", |
| "intensity", |
| "rate_bucket", |
| "loudness_bucket", |
| "group", |
| "split", |
| ) |
|
|
| Float32Array = npt.NDArray[np.float32] |
|
|
|
|
| class DataFormatError(ValueError): |
| """The latent dataset on disk violates the documented format.""" |
|
|
|
|
| @dataclass(frozen=True) |
| class Segment: |
| """One training segment. ``text_ids`` are filled in by the loader (empty before).""" |
|
|
| segment_id: str |
| shard: str |
| offset: int |
| num_frames: int |
| text: str |
| emotion: str | None |
| intensity: float | None |
| rate_bucket: int | None |
| loudness_bucket: int | None |
| group: str |
| split: str |
| text_ids: tuple[int, ...] = () |
| language: str = "tr" |
| |
| |
| units: tuple[tuple[int, int, int], ...] = () |
|
|
|
|
| @dataclass(frozen=True) |
| class Rejection: |
| """An item that was loaded from the index but not admitted to the dataset.""" |
|
|
| segment_id: str |
| reason: str |
| detail: str |
|
|
|
|
| @dataclass |
| class LoadReport: |
| """What the loader did with the index (every record is accounted for).""" |
|
|
| data_dir: str |
| split: str |
| num_index_records: int = 0 |
| num_in_split: int = 0 |
| num_accepted: int = 0 |
| accepted_frames: int = 0 |
| latent_dim: int = 0 |
| rejection_counts: dict[str, int] = field(default_factory=dict) |
| rejections: list[Rejection] = field(default_factory=list) |
| |
| per_dir: list[LoadReport] = field(default_factory=list) |
|
|
| @property |
| def num_rejected(self) -> int: |
| return sum(self.rejection_counts.values()) |
|
|
| def to_dict(self, *, max_examples: int = 20) -> dict[str, object]: |
| """JSON summary (the first ``max_examples`` rejections as examples). |
| |
| A merged multi-directory report adds a ``per_dir`` list (one summary per directory, |
| in load order) and joins the directories with ``"; "`` in ``data_dir``; the key is |
| absent for a single directory, whose summary is unchanged. |
| """ |
| summary: dict[str, object] = { |
| "data_dir": self.data_dir, |
| "split": self.split, |
| "num_index_records": self.num_index_records, |
| "num_in_split": self.num_in_split, |
| "num_accepted": self.num_accepted, |
| "num_rejected": self.num_rejected, |
| "accepted_frames": self.accepted_frames, |
| "latent_dim": self.latent_dim, |
| "rejection_counts": dict(sorted(self.rejection_counts.items())), |
| "rejection_examples": [ |
| {"segment_id": r.segment_id, "reason": r.reason, "detail": r.detail} |
| for r in self.rejections[:max_examples] |
| ], |
| } |
| if self.per_dir: |
| summary["per_dir"] = [r.to_dict(max_examples=max_examples) for r in self.per_dir] |
| return summary |
|
|
|
|
| def _fail(line_no: int, message: str) -> DataFormatError: |
| return DataFormatError(f"{INDEX_NAME} line {line_no}: {message}") |
|
|
|
|
| def _is_int(value: object) -> bool: |
| return isinstance(value, int) and not isinstance(value, bool) |
|
|
|
|
| def _parse_record(raw: object, line_no: int) -> Segment: |
| """Validate one decoded JSON object and build a :class:`Segment` (no text ids yet).""" |
| if not isinstance(raw, dict): |
| raise _fail(line_no, f"expected a JSON object, got {type(raw).__name__}") |
| keys = set(raw) - {"language"} |
| if keys != set(_RECORD_KEYS): |
| raise _fail( |
| line_no, |
| f"keys differ: missing {sorted(set(_RECORD_KEYS) - keys)}, " |
| f"unknown {sorted(keys - set(_RECORD_KEYS))}", |
| ) |
| segment_id, shard, text, group, split = ( |
| raw["segment_id"], |
| raw["shard"], |
| raw["text"], |
| raw["group"], |
| raw["split"], |
| ) |
| for name, value in ( |
| ("segment_id", segment_id), |
| ("shard", shard), |
| ("text", text), |
| ("group", group), |
| ): |
| if not isinstance(value, str) or not value: |
| raise _fail(line_no, f"{name} must be a non-empty string, got {value!r}") |
| if not _SHARD_RE.fullmatch(shard): |
| raise _fail(line_no, f"shard must be a plain *.npy file name, got {shard!r}") |
| if split not in SPLITS: |
| raise _fail(line_no, f"split must be one of {SPLITS}, got {split!r}") |
| offset, num_frames = raw["offset"], raw["num_frames"] |
| if not _is_int(offset) or offset < 0: |
| raise _fail(line_no, f"offset must be an int >= 0, got {offset!r}") |
| if not _is_int(num_frames) or num_frames < 1: |
| raise _fail(line_no, f"num_frames must be an int >= 1, got {num_frames!r}") |
| emotion = raw["emotion"] |
| if emotion is not None and emotion not in EMOTIONS: |
| raise _fail(line_no, f"emotion must be null or one of {EMOTIONS}, got {emotion!r}") |
| intensity = raw["intensity"] |
| if intensity is not None: |
| if isinstance(intensity, bool) or not isinstance(intensity, int | float): |
| raise _fail(line_no, f"intensity must be null or a number, got {intensity!r}") |
| intensity = float(intensity) |
| if not 0.0 <= intensity <= 1.0: |
| raise _fail(line_no, f"intensity must be in [0, 1], got {intensity!r}") |
| for name, upper in (("rate_bucket", N_RATE_BUCKETS), ("loudness_bucket", N_LOUDNESS_BUCKETS)): |
| value = raw[name] |
| if value is not None and (not _is_int(value) or not 0 <= value < upper): |
| raise _fail(line_no, f"{name} must be null or an int in 0..{upper - 1}, got {value!r}") |
| language = raw.get("language", "tr") |
| if language not in LANGUAGES: |
| raise _fail(line_no, f"language must be one of {LANGUAGES}, got {language!r}") |
| return Segment( |
| segment_id=segment_id, |
| shard=shard, |
| offset=offset, |
| num_frames=num_frames, |
| text=text, |
| emotion=emotion, |
| intensity=intensity, |
| rate_bucket=raw["rate_bucket"], |
| loudness_bucket=raw["loudness_bucket"], |
| group=group, |
| split=split, |
| language=language, |
| ) |
|
|
|
|
| def _read_index(path: Path) -> list[Segment]: |
| if not path.is_file(): |
| raise DataFormatError(f"missing {path}") |
| records: list[Segment] = [] |
| seen: set[str] = set() |
| with path.open(encoding="utf-8") as handle: |
| for line_no, line in enumerate(handle, start=1): |
| if not line.strip(): |
| continue |
| try: |
| raw = json.loads(line) |
| except json.JSONDecodeError as exc: |
| raise _fail(line_no, f"invalid JSON ({exc})") from exc |
| record = _parse_record(raw, line_no) |
| if record.segment_id in seen: |
| raise _fail(line_no, f"duplicate segment_id {record.segment_id!r}") |
| seen.add(record.segment_id) |
| records.append(record) |
| return records |
|
|
|
|
| class LatentDataset: |
| """Accepted segments of one split, with memory-mapped latent shards. |
| |
| ``data_dir`` is the (first) directory; ``data_dirs`` lists every directory of a merged |
| multi-directory dataset (see :func:`load_latent_datasets`), whose ``Segment.shard`` and |
| shard keys are then ``"<dir index>/<file name>"``. |
| |
| Not thread-safe to :meth:`close` while in use. ``close()`` drops the memory maps (needed |
| on Windows before deleting the files). |
| """ |
|
|
| def __init__( |
| self, |
| data_dir: Path, |
| segments: Sequence[Segment], |
| shards: dict[str, npt.NDArray[np.floating]], |
| latent_dim: int, |
| *, |
| data_dirs: Sequence[Path] | None = None, |
| ) -> None: |
| self.data_dir = data_dir |
| self.data_dirs: tuple[Path, ...] = (data_dir,) if data_dirs is None else tuple(data_dirs) |
| self.segments: tuple[Segment, ...] = tuple(segments) |
| self._shards = shards |
| self.latent_dim = latent_dim |
|
|
| def __len__(self) -> int: |
| return len(self.segments) |
|
|
| def __getitem__(self, index: int) -> Segment: |
| return self.segments[index] |
|
|
| @property |
| def lengths(self) -> list[int]: |
| """``num_frames`` per segment, in dataset order.""" |
| return [s.num_frames for s in self.segments] |
|
|
| def load_latents(self, index: int) -> Float32Array: |
| """The segment's latent ``[num_frames, latent_dim]`` as a float32 copy. |
| |
| Raises: |
| DataFormatError: the stored values contain NaN/Inf. |
| """ |
| seg = self.segments[index] |
| rows = self._shards[seg.shard][seg.offset : seg.offset + seg.num_frames] |
| out: Float32Array = np.array(rows, dtype=np.float32) |
| if not bool(np.isfinite(out).all()): |
| raise DataFormatError(f"non-finite latent values in segment {seg.segment_id!r}") |
| return out |
|
|
| def fingerprint(self) -> str: |
| """SHA-256 over (id, shard, offset, frames) of every accepted segment, in order. |
| |
| Stored in checkpoints: a resume against different data (or a different rejection |
| outcome) must fail, since the sampler position would be meaningless. |
| """ |
| digest = hashlib.sha256() |
| for s in self.segments: |
| digest.update(f"{s.segment_id}\t{s.shard}\t{s.offset}\t{s.num_frames}\n".encode()) |
| if s.units: |
| digest.update(f"units\t{s.units}\n".encode()) |
| return digest.hexdigest() |
|
|
| def close(self) -> None: |
| self._shards.clear() |
|
|
|
|
| |
| MIN_LATENT_STD = 1e-3 |
|
|
|
|
| @dataclass(frozen=True) |
| class LatentStats: |
| """Per-channel mean / std of the training latents (the model sees ``(x - mean) / std``). |
| |
| Flow matching draws ``x0 ~ N(0, I)``; raw VAE latents have arbitrary per-channel offset |
| and scale, so without this the noise and data ends of the path are mismatched. The |
| stats belong to the model: they are stored in every checkpoint and are needed to turn |
| generated latents back into VAE latents (:meth:`denormalize`). |
| """ |
|
|
| mean: tuple[float, ...] |
| std: tuple[float, ...] |
|
|
| def __post_init__(self) -> None: |
| if not self.mean or len(self.mean) != len(self.std): |
| raise ValueError("mean and std must be non-empty and of equal length") |
| if not all(np.isfinite(self.mean)) or not all( |
| np.isfinite(s) and s >= MIN_LATENT_STD for s in self.std |
| ): |
| raise ValueError(f"latent stats must be finite with std >= {MIN_LATENT_STD}") |
|
|
| def close_to(self, other: LatentStats, rtol: float = 1e-6) -> bool: |
| """Same channels and values equal up to float summation noise.""" |
| return len(self.mean) == len(other.mean) and bool( |
| np.allclose(self.mean, other.mean, rtol=rtol, atol=1e-9) |
| and np.allclose(self.std, other.std, rtol=rtol, atol=1e-9) |
| ) |
|
|
| def normalize(self, latents: Float32Array) -> Float32Array: |
| """``latents: (num_frames, latent_dim)`` -> standardised float32 copy.""" |
| mean = np.asarray(self.mean, dtype=np.float32) |
| std = np.asarray(self.std, dtype=np.float32) |
| out: Float32Array = ((latents - mean) / std).astype(np.float32) |
| return out |
|
|
| def denormalize(self, latents: Float32Array) -> Float32Array: |
| """Inverse of :meth:`normalize`.""" |
| mean = np.asarray(self.mean, dtype=np.float32) |
| std = np.asarray(self.std, dtype=np.float32) |
| out: Float32Array = (latents * std + mean).astype(np.float32) |
| return out |
|
|
| def to_dict(self) -> dict[str, list[float]]: |
| return {"mean": list(self.mean), "std": list(self.std)} |
|
|
| @classmethod |
| def from_mapping(cls, raw: object) -> LatentStats: |
| if not isinstance(raw, dict) or set(raw) != {"mean", "std"}: |
| raise ValueError(f"latent stats must be a mapping with mean/std, got {raw!r}") |
| return cls(tuple(float(x) for x in raw["mean"]), tuple(float(x) for x in raw["std"])) |
|
|
|
|
| def compute_latent_stats(data: LatentDataset) -> LatentStats: |
| """Per-channel mean/std over every frame of ``data`` (float64 accumulation). |
| |
| ``std`` is floored at :data:`MIN_LATENT_STD`. |
| |
| Raises: |
| ValueError: ``data`` is empty. |
| """ |
| if len(data) == 0: |
| raise ValueError("cannot compute latent stats of an empty dataset") |
| total = np.zeros(data.latent_dim, dtype=np.float64) |
| total_sq = np.zeros(data.latent_dim, dtype=np.float64) |
| count = 0 |
| for index in range(len(data)): |
| lat = data.load_latents(index).astype(np.float64) |
| total += lat.sum(axis=0) |
| total_sq += np.square(lat).sum(axis=0) |
| count += lat.shape[0] |
| mean = total / count |
| var = np.maximum(total_sq / count - np.square(mean), 0.0) |
| std = np.maximum(np.sqrt(var), MIN_LATENT_STD) |
| return LatentStats(tuple(float(x) for x in mean), tuple(float(x) for x in std)) |
|
|
|
|
| def _open_shards( |
| data_dir: Path, records: Sequence[Segment] |
| ) -> tuple[dict[str, npt.NDArray[np.floating]], int]: |
| """Memory-map every referenced shard and validate bounds. Returns (shards, latent_dim).""" |
| shards: dict[str, npt.NDArray[np.floating]] = {} |
| latent_dim: int | None = None |
| for name in sorted({r.shard for r in records}): |
| path = data_dir / name |
| if not path.is_file(): |
| raise DataFormatError(f"shard {name!r} referenced by the index is missing") |
| try: |
| arr = np.load(path, mmap_mode="r", allow_pickle=False) |
| except (ValueError, OSError) as exc: |
| raise DataFormatError(f"cannot read shard {name!r}: {exc}") from exc |
| if arr.ndim != 2 or arr.dtype not in (np.float16, np.float32): |
| raise DataFormatError( |
| f"shard {name!r} must be a 2-D float16/float32 array, " |
| f"got dtype {arr.dtype} shape {arr.shape}" |
| ) |
| if latent_dim is None: |
| latent_dim = int(arr.shape[1]) |
| elif int(arr.shape[1]) != latent_dim: |
| raise DataFormatError( |
| f"shard {name!r} has latent_dim {arr.shape[1]}, expected {latent_dim}" |
| ) |
| shards[name] = arr |
| for r in records: |
| rows = shards[r.shard].shape[0] |
| if r.offset + r.num_frames > rows: |
| raise DataFormatError( |
| f"segment {r.segment_id!r}: rows {r.offset}..{r.offset + r.num_frames} " |
| f"exceed shard {r.shard!r} ({rows} rows)" |
| ) |
| return shards, latent_dim if latent_dim is not None else 0 |
|
|
|
|
| def _encode(segment: Segment) -> tuple[Segment | None, Rejection | None]: |
| """Encode the text; return the segment with ids or the rejection.""" |
| try: |
| ids = encode_text(segment.text) |
| except ValueError as exc: |
| return None, Rejection(segment.segment_id, REJECT_TEXT_ENCODING, str(exc)) |
| if not ids: |
| return None, Rejection(segment.segment_id, REJECT_EMPTY_TEXT, "text encodes to no tokens") |
| if len(ids) > segment.num_frames: |
| detail = f"{len(ids)} text tokens > {segment.num_frames} frames" |
| return None, Rejection(segment.segment_id, REJECT_TEXT_TOO_LONG, detail) |
| return replace(segment, text_ids=tuple(ids)), None |
|
|
|
|
| def _check_groups_disjoint(records: Sequence[Segment]) -> None: |
| """A group (speaker / source video) in two splits would leak voices into evaluation.""" |
| splits_of: dict[str, set[str]] = {} |
| for record in records: |
| splits_of.setdefault(record.group, set()).add(record.split) |
| crossing = sorted(g for g, s in splits_of.items() if len(s) > 1) |
| if crossing: |
| raise DataFormatError( |
| f"{len(crossing)} group(s) appear in more than one split, e.g. " |
| f"{crossing[0]!r} in {sorted(splits_of[crossing[0]])}" |
| ) |
|
|
|
|
| def _attach_timings( |
| segment: Segment, timings: Mapping[str, SegmentTimings] |
| ) -> tuple[Segment | None, Rejection | None]: |
| """E11: attach the unit timeline; reject a segment without one or with a mismatch.""" |
| found = timings.get(segment.segment_id) |
| if found is None: |
| return None, Rejection(segment.segment_id, REJECT_NO_TIMINGS, "not in timings.jsonl") |
| if found.n_tokens != len(segment.text_ids) or found.num_frames != segment.num_frames: |
| detail = ( |
| f"timings {found.n_tokens} tokens / {found.num_frames} frames != " |
| f"{len(segment.text_ids)} tokens / {segment.num_frames} frames" |
| ) |
| return None, Rejection(segment.segment_id, REJECT_TIMINGS_MISMATCH, detail) |
| units = tuple((u.char_start, u.char_end, u.frame_start) for u in found.units) |
| return replace(segment, units=units), None |
|
|
|
|
| def _load_split( |
| data_dir: Path, split: str, records: Sequence[Segment], *, timings: bool = False |
| ) -> tuple[LatentDataset, LoadReport]: |
| """Open, validate and encode the ``split`` records of one directory's (checked) index. |
| |
| With ``timings`` (aligned text path) every accepted segment also gets its unit |
| timeline from ``timings.jsonl`` (:func:`nida.data.timings.load_timings`). |
| """ |
| unit_map = load_timings(data_dir) if timings else None |
| in_split = [r for r in records if r.split == split] |
| shards, latent_dim = _open_shards(data_dir, in_split) |
| report = LoadReport( |
| data_dir=str(data_dir), |
| split=split, |
| num_index_records=len(records), |
| num_in_split=len(in_split), |
| latent_dim=latent_dim, |
| ) |
| accepted: list[Segment] = [] |
| counts: Counter[str] = Counter() |
| for record in in_split: |
| encoded, rejection = _encode(record) |
| if encoded is not None and unit_map is not None: |
| encoded, rejection = _attach_timings(encoded, unit_map) |
| if encoded is not None: |
| accepted.append(encoded) |
| elif rejection is not None: |
| counts[rejection.reason] += 1 |
| report.rejections.append(rejection) |
| report.rejection_counts = dict(counts) |
| report.num_accepted = len(accepted) |
| report.accepted_frames = sum(s.num_frames for s in accepted) |
| if report.num_accepted + report.num_rejected != report.num_in_split: |
| raise AssertionError("loader lost track of records") |
| return LatentDataset(data_dir, accepted, shards, latent_dim), report |
|
|
|
|
| def load_latent_dataset( |
| data_dir: Path, split: str, *, timings: bool = False |
| ) -> tuple[LatentDataset, LoadReport]: |
| """Load one split of a latent dataset directory (format: see module docstring). |
| |
| Returns the dataset (accepted items, in index order) and the report that accounts for |
| every record of the split. The caller decides whether an empty dataset is fatal. |
| |
| Raises: |
| DataFormatError: structural problems (see module docstring). |
| ValueError: unknown ``split``. |
| """ |
| if split not in SPLITS: |
| raise ValueError(f"split must be one of {SPLITS}, got {split!r}") |
| records = _read_index(data_dir / INDEX_NAME) |
| _check_groups_disjoint(records) |
| return _load_split(data_dir, split, records, timings=timings) |
|
|
|
|
| def _check_cross_directory(data_dirs: Sequence[Path], per_dir: Sequence[Sequence[Segment]]) -> None: |
| """Checks that only make sense across directories (each index is already valid alone).""" |
| owner: dict[str, int] = {} |
| for k, records in enumerate(per_dir): |
| for record in records: |
| first = owner.setdefault(record.segment_id, k) |
| if first != k: |
| raise DataFormatError( |
| f"segment_id {record.segment_id!r} appears in both {data_dirs[first]} " |
| f"and {data_dirs[k]}" |
| ) |
| try: |
| _check_groups_disjoint([r for records in per_dir for r in records]) |
| except DataFormatError as exc: |
| raise DataFormatError( |
| f"across data directories {[str(d) for d in data_dirs]}: {exc}" |
| ) from exc |
|
|
|
|
| |
| |
| |
| META_MUST_MATCH = ( |
| "frame_hop_ms", |
| "rate_bucket_edges_chars_per_s", |
| "loudness_bucket_edges_dbfs", |
| ) |
| META_NAME = "meta.json" |
|
|
|
|
| def _check_meta_compatible(data_dirs: Sequence[Path]) -> None: |
| """Directories with a ``meta.json`` must agree on :data:`META_MUST_MATCH`; a mix of |
| directories with and without one is refused (the missing one cannot be checked).""" |
| metas: list[tuple[Path, Mapping[str, object]]] = [] |
| for data_dir in data_dirs: |
| path = data_dir / META_NAME |
| if path.is_file(): |
| try: |
| loaded = json.loads(path.read_text(encoding="utf-8")) |
| except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc: |
| raise DataFormatError(f"{path}: unreadable meta.json: {exc}") from exc |
| if not isinstance(loaded, dict): |
| raise DataFormatError(f"{path}: meta.json must be an object") |
| metas.append((data_dir, loaded)) |
| if not metas: |
| return |
| if len(metas) != len(data_dirs): |
| missing = [str(d) for d in data_dirs if not (d / META_NAME).is_file()] |
| raise DataFormatError(f"cannot merge: no {META_NAME} in {missing} (others have one)") |
| first_dir, first = metas[0] |
| for data_dir, meta in metas[1:]: |
| for key in META_MUST_MATCH: |
| if meta.get(key) != first.get(key): |
| raise DataFormatError( |
| f"cannot merge: {key} differs ({first_dir}: {first.get(key)!r}, " |
| f"{data_dir}: {meta.get(key)!r})" |
| ) |
|
|
|
|
| def _check_latent_dims(data_dirs: Sequence[Path], reports: Sequence[LoadReport]) -> None: |
| known = [(d, r.latent_dim) for d, r in zip(data_dirs, reports, strict=True) if r.latent_dim] |
| for data_dir, dim in known[1:]: |
| if dim != known[0][1]: |
| raise DataFormatError( |
| f"latent_dim differs across data directories: {known[0][0]} has {known[0][1]}, " |
| f"{data_dir} has {dim}" |
| ) |
|
|
|
|
| def _merge_reports(split: str, reports: Sequence[LoadReport]) -> LoadReport: |
| counts: Counter[str] = Counter() |
| for report in reports: |
| counts.update(report.rejection_counts) |
| return LoadReport( |
| data_dir="; ".join(r.data_dir for r in reports), |
| split=split, |
| num_index_records=sum(r.num_index_records for r in reports), |
| num_in_split=sum(r.num_in_split for r in reports), |
| num_accepted=sum(r.num_accepted for r in reports), |
| accepted_frames=sum(r.accepted_frames for r in reports), |
| latent_dim=max(r.latent_dim for r in reports), |
| rejection_counts=dict(counts), |
| rejections=[x for r in reports for x in r.rejections], |
| per_dir=list(reports), |
| ) |
|
|
|
|
| def load_latent_datasets( |
| data_dirs: Sequence[Path], split: str, *, timings: bool = False |
| ) -> tuple[LatentDataset, LoadReport]: |
| """Load one split of several latent dataset directories as one union dataset. |
| |
| Segments keep directory order, then index order within a directory. With a single |
| directory this is exactly :func:`load_latent_dataset` (same segments, same |
| ``fingerprint()``, so existing checkpoints still resume). With several, every shard key |
| and ``Segment.shard`` is prefixed with ``"<dir index>/"`` to keep shards of different |
| directories apart (the fingerprint therefore depends on the directory order), and the |
| returned report sums the counts and carries the per-directory reports in ``per_dir``. |
| |
| Cross-directory checks (the index records of *all* splits are considered, not just |
| ``split``): a ``segment_id`` in two directories, a ``group`` that is in different |
| splits in different directories, and differing ``latent_dim``. A directory with no |
| records in ``split`` has no shard to read, so its ``latent_dim`` is unknown and is not |
| compared. |
| |
| Raises: |
| DataFormatError: structural problems in any directory, or a failed cross check. |
| ValueError: unknown ``split``, no directories, or the same directory given twice. |
| """ |
| if split not in SPLITS: |
| raise ValueError(f"split must be one of {SPLITS}, got {split!r}") |
| if not data_dirs: |
| raise ValueError("at least one latent dataset directory is required") |
| if len(data_dirs) == 1: |
| return load_latent_dataset(data_dirs[0], split, timings=timings) |
| resolved = [d.resolve() for d in data_dirs] |
| if len(set(resolved)) != len(resolved): |
| raise ValueError(f"the same latent dataset directory is given twice: {list(data_dirs)}") |
| per_dir_records: list[list[Segment]] = [] |
| for data_dir in data_dirs: |
| records = _read_index(data_dir / INDEX_NAME) |
| _check_groups_disjoint(records) |
| per_dir_records.append(records) |
| _check_cross_directory(data_dirs, per_dir_records) |
| _check_meta_compatible(data_dirs) |
|
|
| loaded: list[tuple[LatentDataset, LoadReport]] = [] |
| try: |
| for data_dir, records in zip(data_dirs, per_dir_records, strict=True): |
| loaded.append(_load_split(data_dir, split, records, timings=timings)) |
| _check_latent_dims(data_dirs, [report for _, report in loaded]) |
| except BaseException: |
| for dataset, _ in loaded: |
| dataset.close() |
| raise |
| segments: list[Segment] = [] |
| shards: dict[str, npt.NDArray[np.floating]] = {} |
| for k, (dataset, _) in enumerate(loaded): |
| segments.extend(replace(s, shard=f"{k}/{s.shard}") for s in dataset.segments) |
| shards.update({f"{k}/{name}": arr for name, arr in dataset._shards.items()}) |
| report = _merge_reports(split, [report for _, report in loaded]) |
| merged = LatentDataset( |
| data_dirs[0], segments, shards, report.latent_dim, data_dirs=list(data_dirs) |
| ) |
| return merged, report |
|
|
