src/turn_detection/training/datasets.py

"""Flexible audio/feature datasets used by the command-line training tools.

The code accepts JSONL manifests, precomputed feature records, or any source
understood by :func:`turn_detection.data.iter_records`.  Audio is decoded lazily
so a 41 GB corpus is never materialized in Python memory.
"""

from __future__ import annotations

import io
import json
import random
import wave
from collections import deque
from collections.abc import Iterable, Iterator, Mapping, Sequence
from pathlib import Path
from typing import Any

import torch
from torch import Tensor
from torch.utils.data import DataLoader, Dataset, IterableDataset, get_worker_info

from turn_detection.models.features import LogMelFrontend


def _field(record: Any, *names: str, default: Any = None) -> Any:
    for name in names:
        if isinstance(record, Mapping) and name in record:
            return record[name]
        if hasattr(record, name):
            return getattr(record, name)
    return default


def _read_jsonl(path: Path, split: str | None) -> list[dict[str, Any]]:
    records: list[dict[str, Any]] = []
    with path.open("r", encoding="utf-8") as handle:
        for line_number, line in enumerate(handle, start=1):
            if not line.strip():
                continue
            value = json.loads(line)
            if not isinstance(value, dict):
                raise ValueError(f"{path}:{line_number}: expected an object")
            if split is not None and value.get("split") not in (None, split):
                continue
            value["__manifest_dir"] = str(path.parent)
            records.append(value)
    return records


class ManifestDataset(Dataset):
    def __init__(self, records: Sequence[Any]) -> None:
        self.records = list(records)

    def __len__(self) -> int:
        return len(self.records)

    def __getitem__(self, index: int) -> Any:
        return self.records[index]


class RecordStream(IterableDataset):
    """Worker-sharded stream with deterministic bounded-buffer shuffling."""

    def __init__(
        self,
        source: str,
        split: str,
        revision: str | None = None,
        token: str | None = None,
        shuffle_buffer: int = 2_048,
        seed: int = 17,
        max_examples: int | None = None,
    ) -> None:
        super().__init__()
        self.source = source
        self.split = split
        self.revision = revision
        self.token = token
        self.shuffle_buffer = max(1, min(int(shuffle_buffer), 64))
        self.seed = seed
        self.max_examples = max_examples
        self.epoch = 0

    def set_epoch(self, epoch: int) -> None:
        self.epoch = int(epoch)

    def _records(self) -> Iterable[Any]:
        from turn_detection.data import iter_records

        return iter_records(
            self.source,
            split=self.split,
            revision=self.revision,
            token=self.token,
        )

    def __iter__(self) -> Iterator[Any]:
        worker = get_worker_info()
        worker_id = 0 if worker is None else worker.id
        worker_count = 1 if worker is None else worker.num_workers
        rng = random.Random(self.seed + worker_id + self.epoch * 1_000_003)
        buffer: list[Any] = []
        emitted = 0
        worker_limit = (
            None
            if self.max_examples is None
            else (self.max_examples + worker_count - 1) // worker_count
        )
        for index, record in enumerate(self._records()):
            if index % worker_count != worker_id:
                continue
            if worker_limit is not None and emitted >= worker_limit:
                break
            if self.shuffle_buffer <= 1:
                emitted += 1
                yield record
                continue
            if len(buffer) < self.shuffle_buffer:
                buffer.append(record)
                continue
            selected = rng.randrange(len(buffer))
            emitted += 1
            yield buffer[selected]
            buffer[selected] = record
        rng.shuffle(buffer)
        for record in buffer:
            if worker_limit is not None and emitted >= worker_limit:
                break
            emitted += 1
            yield record


class ManifestAudioStream(IterableDataset):
    """Resolve lightweight manifest rows to Parquet audio in source order."""

    def __init__(
        self,
        rows: Sequence[Mapping[str, Any]],
        source_root: str | Path,
        shuffle_buffer: int = 32,
        seed: int = 17,
        max_examples: int | None = None,
    ) -> None:
        super().__init__()
        self.rows = list(rows)
        self.source_root = str(source_root)
        # Decoded audio can be large. Bound this independently of a caller's
        # metadata shuffle setting to avoid multi-gigabyte worker buffers.
        self.shuffle_buffer = max(1, min(int(shuffle_buffer), 64))
        self.seed = seed
        self.max_examples = max_examples
        self.epoch = 0

    def set_epoch(self, epoch: int) -> None:
        self.epoch = int(epoch)

    def __iter__(self) -> Iterator[Any]:
        from turn_detection.data import iter_manifest_records

        worker = get_worker_info()
        worker_id = 0 if worker is None else worker.id
        worker_count = 1 if worker is None else worker.num_workers
        selected_rows = (
            row for index, row in enumerate(self.rows) if index % worker_count == worker_id
        )
        pending_rows: deque[Mapping[str, Any]] = deque()

        def tracked_rows() -> Iterator[Mapping[str, Any]]:
            for row in selected_rows:
                pending_rows.append(row)
                yield row

        resolved = iter_manifest_records(
            tracked_rows(),
            source_root=self.source_root,
            max_cached_row_groups=1,
        )
        rng = random.Random(self.seed + worker_id + self.epoch * 1_000_003)
        buffer: list[Any] = []
        emitted = 0
        worker_limit = (
            None
            if self.max_examples is None
            else (self.max_examples + worker_count - 1) // worker_count
        )
        for source_record in resolved:
            # Preserve audit/split metadata (duration, group, record ID) while
            # allowing the raw source's audio payload and labels to win.
            record = {**pending_rows.popleft(), **source_record}
            if worker_limit is not None and emitted >= worker_limit:
                break
            if self.shuffle_buffer <= 1:
                emitted += 1
                yield record
            elif len(buffer) < self.shuffle_buffer:
                buffer.append(record)
            else:
                selected = rng.randrange(len(buffer))
                emitted += 1
                yield buffer[selected]
                buffer[selected] = record
        rng.shuffle(buffer)
        for record in buffer:
            if worker_limit is not None and emitted >= worker_limit:
                break
            emitted += 1
            yield record


def _decode_pcm_wav(payload: bytes) -> tuple[Tensor, int]:
    """stdlib fallback for ordinary PCM WAV when soundfile is unavailable."""

    with wave.open(io.BytesIO(payload), "rb") as handle:
        channels = handle.getnchannels()
        sample_rate = handle.getframerate()
        width = handle.getsampwidth()
        frames = handle.readframes(handle.getnframes())
    if width not in (1, 2, 4):
        raise ValueError(f"unsupported PCM sample width: {width}")
    try:
        import numpy as np
    except ImportError as exc:
        raise ImportError("decoding PCM audio requires numpy or soundfile") from exc
    dtype = {1: np.uint8, 2: np.dtype("<i2"), 4: np.dtype("<i4")}[width]
    array = np.frombuffer(frames, dtype=dtype).reshape(-1, channels)
    if width == 1:
        array = (array.astype("float32") - 128.0) / 128.0
    else:
        array = array.astype("float32") / float(2 ** (width * 8 - 1))
    return torch.from_numpy(array).mean(dim=1), sample_rate


def _decode_with_soundfile(source: Any) -> tuple[Tensor, int]:
    try:
        import soundfile as sf
    except ImportError:
        if isinstance(source, bytes | bytearray):
            return _decode_pcm_wav(bytes(source))
        return _decode_pcm_wav(Path(source).read_bytes())
    array, sample_rate = sf.read(source, dtype="float32", always_2d=True)
    return torch.from_numpy(array).mean(dim=1), int(sample_rate)


def _as_mono_float(value: Any) -> Tensor:
    waveform = torch.as_tensor(value)
    if not waveform.dtype.is_floating_point:
        info = torch.iinfo(waveform.dtype)
        maximum = float(max(abs(info.min), info.max))
        waveform = waveform.to(torch.float32) / maximum
    else:
        waveform = waveform.to(torch.float32)
    waveform = waveform.squeeze()
    if waveform.ndim == 2:
        # Audio libraries use either [channels, samples] or [samples, channels].
        channel_axis = 0 if waveform.shape[0] <= 8 else 1
        waveform = waveform.mean(dim=channel_axis)
    if waveform.ndim != 1:
        raise ValueError("decoded audio must be mono or two-dimensional")
    if waveform.numel() == 0:
        raise ValueError("audio cannot be empty")
    waveform = torch.nan_to_num(waveform, nan=0.0, posinf=1.0, neginf=-1.0)
    peak = float(waveform.abs().max()) if waveform.numel() else 0.0
    if peak > 1.0:
        waveform = waveform / peak
    return waveform.clamp(-1.0, 1.0)


def linear_resample_waveform(
    waveform: Tensor, source_sample_rate: int, target_sample_rate: int
) -> Tensor:
    """Deterministic resampling contract shared with the dependency-light runtime."""

    if source_sample_rate <= 0 or target_sample_rate <= 0:
        raise ValueError("sample rates must be positive")
    if source_sample_rate == target_sample_rate:
        return waveform.to(torch.float32)
    try:
        import numpy as np
    except ImportError as exc:  # pragma: no cover - numpy is a base dependency
        raise ImportError("audio resampling requires numpy") from exc
    samples = waveform.detach().cpu().numpy().astype("float32", copy=False)
    output_length = max(1, round(len(samples) * target_sample_rate / source_sample_rate))
    old_x = np.linspace(0.0, 1.0, len(samples), endpoint=False)
    new_x = np.linspace(0.0, 1.0, output_length, endpoint=False)
    output = np.interp(new_x, old_x, samples).astype("float32")
    return torch.from_numpy(output)


def decode_audio(record: Any, target_sample_rate: int, max_seconds: float) -> Tensor:
    """Decode common HF, torchcodec, file, byte, and array audio representations."""

    audio = _field(record, "audio")
    sample_rate = _field(record, "sample_rate", default=None)
    waveform: Tensor | None = None

    if audio is not None and hasattr(audio, "get_all_samples"):
        samples = audio.get_all_samples()
        waveform = _as_mono_float(samples.data)
        sample_rate = int(samples.sample_rate)
    elif isinstance(audio, Mapping):
        sample_rate = audio.get("sampling_rate", audio.get("sample_rate", sample_rate))
        if audio.get("array") is not None:
            waveform = _as_mono_float(audio["array"])
        elif audio.get("samples") is not None:
            waveform = _as_mono_float(audio["samples"])
        elif audio.get("bytes") is not None:
            waveform, decoded_rate = _decode_with_soundfile(io.BytesIO(audio["bytes"]))
            sample_rate = sample_rate or decoded_rate
        elif audio.get("path"):
            waveform, decoded_rate = _decode_with_soundfile(audio["path"])
            sample_rate = sample_rate or decoded_rate
    elif (
        isinstance(audio, Tensor | list | tuple)
        or audio is not None
        and hasattr(audio, "__array__")
    ):
        waveform = _as_mono_float(audio)
    elif isinstance(audio, bytes | bytearray):
        waveform, decoded_rate = _decode_with_soundfile(io.BytesIO(audio))
        sample_rate = sample_rate or decoded_rate

    if waveform is None:
        path_value = _field(record, "audio_path", "path")
        if not path_value:
            raise ValueError(
                "record contains neither decodable audio nor audio_path; use raw records "
                "or generate a manifest with extracted audio paths"
            )
        path = Path(path_value)
        manifest_dir = _field(record, "__manifest_dir")
        if not path.is_absolute() and manifest_dir:
            path = Path(manifest_dir) / path
        waveform, decoded_rate = _decode_with_soundfile(path)
        sample_rate = sample_rate or decoded_rate

    if sample_rate is None:
        raise ValueError("audio sample rate is missing")
    sample_rate = int(sample_rate)
    if sample_rate <= 0:
        raise ValueError("audio sample rate must be positive")
    source_maximum_samples = int(round(max_seconds * sample_rate))
    if source_maximum_samples <= 0:
        raise ValueError("max_seconds must be positive")
    # Crop before resampling to bound memory/compute for very long recordings;
    # the dependency-light runtime applies this exact suffix rule too.
    waveform = waveform[-source_maximum_samples:]
    if sample_rate != target_sample_rate:
        waveform = linear_resample_waveform(waveform, sample_rate, target_sample_rate)
    maximum_samples = int(round(max_seconds * target_sample_rate))
    # Turn intent is concentrated at the end, so crop the suffix rather than the prefix.
    return waveform[-maximum_samples:].contiguous()


def _load_precomputed(record: Any) -> Tensor | None:
    value = _field(record, "log_mel", "input_features", "features")
    if value is not None:
        return torch.as_tensor(value, dtype=torch.float32)
    feature_path = _field(record, "feature_path")
    if not feature_path:
        return None
    path = Path(feature_path)
    manifest_dir = _field(record, "__manifest_dir")
    if not path.is_absolute() and manifest_dir:
        path = Path(manifest_dir) / path
    loaded = torch.load(path, map_location="cpu", weights_only=True)
    if isinstance(loaded, Mapping):
        loaded = loaded.get("log_mel", loaded.get("features"))
    if loaded is None:
        raise ValueError(f"feature file {path} does not contain log_mel/features")
    return torch.as_tensor(loaded, dtype=torch.float32)


class AudioFeatureCollator:
    def __init__(
        self,
        frontend: LogMelFrontend,
        max_seconds: float,
    ) -> None:
        self.frontend = frontend.cpu().eval()
        self.max_seconds = max_seconds

    def __call__(self, records: Sequence[Any]) -> dict[str, Any]:
        if not records:
            raise ValueError("cannot collate an empty batch")
        precomputed = [_load_precomputed(record) for record in records]
        if all(value is not None for value in precomputed):
            features = [value for value in precomputed if value is not None]
            for feature in features:
                if feature.ndim != 2:
                    raise ValueError("precomputed features must be [n_mels, frames]")
            max_frames = max(feature.shape[-1] for feature in features)
            batch = features[0].new_zeros((len(features), features[0].shape[0], max_frames))
            frame_mask = torch.zeros((len(features), max_frames), dtype=torch.bool)
            for index, feature in enumerate(features):
                frames = feature.shape[-1]
                batch[index, :, :frames] = feature
                frame_mask[index, :frames] = True
        elif any(value is not None for value in precomputed):
            raise ValueError("a batch cannot mix waveform and precomputed feature records")
        else:
            waveforms = [
                decode_audio(record, self.frontend.config.sample_rate, self.max_seconds)
                for record in records
            ]
            lengths = torch.tensor([waveform.numel() for waveform in waveforms], dtype=torch.long)
            target_samples = int(round(self.max_seconds * self.frontend.config.sample_rate))
            padded = torch.zeros((len(waveforms), target_samples), dtype=torch.float32)
            for index, waveform in enumerate(waveforms):
                if self.frontend.config.pad_side == "left":
                    padded[index, -waveform.numel() :] = waveform
                else:
                    padded[index, : waveform.numel()] = waveform
            with torch.no_grad():
                batch, frame_mask = self.frontend(padded, lengths)

        def label(*names: str, missing: float = -1.0) -> Tensor:
            values = [_field(record, *names, default=missing) for record in records]
            values = [missing if value is None else float(value) for value in values]
            return torch.tensor(values, dtype=torch.float32)

        endpoint_values = [
            _field(record, "endpoint", "endpoint_bool", "label", default=None) for record in records
        ]
        if any(value is None for value in endpoint_values):
            raise ValueError("every training record must have an endpoint label")

        return {
            "log_mel": batch,
            "attention_mask": frame_mask,
            "endpoint": torch.tensor(
                [float(value) for value in endpoint_values], dtype=torch.float32
            ),
            "midfiller": label("midfiller", "midfiller_bool"),
            "endfiller": label("endfiller", "endfiller_bool"),
            "record_id": [str(_field(record, "record_id", "id", default="")) for record in records],
            "turn_id": [
                str(_field(record, "turn_id", "conversation_id", "record_id", "id", default=""))
                for record in records
            ],
            # Keep provenance for operational evaluation. Falling back to a
            # record ID is useful for joins, but it must never make independent
            # clips look like observed turns/conversations.
            "turn_id_observed": [
                _field(record, "turn_id", "conversation_id", default=None) is not None
                for record in records
            ],
            "group_id": [
                str(
                    _field(
                        record,
                        "group_id",
                        "speaker_id",
                        "turn_id",
                        "conversation_id",
                        "record_id",
                        "id",
                        default="",
                    )
                )
                for record in records
            ],
            "language": [_field(record, "language", default="unknown") for record in records],
            "dataset": [
                _field(record, "dataset", "source_dataset", default="unknown") for record in records
            ],
            "synthetic": [_field(record, "synthetic", default=None) for record in records],
            "duration_seconds": [
                _field(record, "duration_seconds", "duration", default=None) for record in records
            ],
        }


class SyntheticFeatureDataset(Dataset):
    """Deterministic learnable smoke dataset; never used for reported results."""

    def __init__(
        self, examples: int = 64, n_mels: int = 80, frames: int = 96, seed: int = 17
    ) -> None:
        generator = torch.Generator().manual_seed(seed)
        self.records: list[dict[str, Any]] = []
        for index in range(examples):
            endpoint = index % 2
            valid_frames = frames - (index % 13)
            features = torch.randn((n_mels, valid_frames), generator=generator) * 0.5
            # A learnable prosodic proxy in the final 12 frames.
            features[:8, -12:] += 1.5 if endpoint else -1.5
            self.records.append(
                {
                    "log_mel": features,
                    "endpoint": endpoint,
                    "midfiller": float(index % 5 == 0) if not endpoint else -1.0,
                    "endfiller": float(index % 7 == 0) if endpoint else -1.0,
                    "record_id": f"synthetic-{index}",
                    "language": "synthetic",
                    "dataset": "smoke",
                    "synthetic": True,
                }
            )

    def __len__(self) -> int:
        return len(self.records)

    def __getitem__(self, index: int) -> Mapping[str, Any]:
        return self.records[index]


def build_record_dataloader(
    source: str | Path,
    split: str,
    frontend: LogMelFrontend,
    batch_size: int,
    max_seconds: float,
    shuffle: bool,
    num_workers: int = 0,
    seed: int = 17,
    revision: str | None = None,
    token: str | None = None,
    shuffle_buffer: int = 2_048,
    max_examples: int | None = None,
    source_root: str | Path | None = None,
) -> DataLoader:
    path = Path(source)
    if path.is_file() and path.suffix.lower() in {".jsonl", ".json"}:
        records = _read_jsonl(path, split)

        def has_resolvable_payload(record: Mapping[str, Any]) -> bool:
            if _field(record, "audio", "feature_path", "log_mel", "features") is not None:
                return True
            audio_path = _field(record, "audio_path")
            if not audio_path:
                return False
            candidate = Path(audio_path)
            if not candidate.is_absolute():
                candidate = path.parent / candidate
            return candidate.is_file()

        sampled = records[: min(128, len(records))]
        has_direct_payload = bool(sampled) and all(
            has_resolvable_payload(record) for record in sampled
        )
        has_source_reference = bool(records) and all(
            _field(record, "source_file") is not None and _field(record, "source_row") is not None
            for record in sampled
        )
        if has_source_reference and not has_direct_payload:
            dataset: Dataset | IterableDataset = ManifestAudioStream(
                records,
                source_root=source_root or path.parent,
                shuffle_buffer=shuffle_buffer if shuffle else 1,
                seed=seed,
                max_examples=max_examples,
            )
            generator = None
            can_shuffle = False
        else:
            if max_examples is not None:
                records = records[:max_examples]
            dataset = ManifestDataset(records)
            generator = torch.Generator().manual_seed(seed)
            can_shuffle = shuffle
    else:
        dataset = RecordStream(
            str(source),
            split=split,
            revision=revision,
            token=token,
            shuffle_buffer=shuffle_buffer if shuffle else 1,
            seed=seed,
            max_examples=max_examples,
        )
        generator = None
        can_shuffle = False
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=can_shuffle,
        num_workers=num_workers,
        collate_fn=AudioFeatureCollator(frontend, max_seconds),
        generator=generator,
        # Recreating workers each epoch propagates IterableDataset.set_epoch.
        persistent_workers=False,
    )


def build_smoke_dataloaders(
    frontend: LogMelFrontend,
    batch_size: int = 8,
    seed: int = 17,
) -> tuple[DataLoader, DataLoader]:
    train = SyntheticFeatureDataset(64, frontend.config.n_mels, 96, seed)
    validation = SyntheticFeatureDataset(32, frontend.config.n_mels, 96, seed + 1)
    collator = AudioFeatureCollator(frontend, max_seconds=1.0)
    generator = torch.Generator().manual_seed(seed)
    return (
        DataLoader(
            train, batch_size=batch_size, shuffle=True, collate_fn=collator, generator=generator
        ),
        DataLoader(validation, batch_size=batch_size, shuffle=False, collate_fn=collator),
    )