diff --git a/README.md b/README.md index 1895568..a1febc5 100644 --- a/README.md +++ b/README.md @@ -66,7 +66,7 @@ shape: (5_271_939, 3) ```pycon >>> x['audio'] -WSAudio(audio_reader=AudioReader(src=, sample_rate=None), tstart=614.46246, tend=627.3976) +WSAudioSegment(episode=WSAudioEpisode(src=, sample_rate=None), tstart=614.46246, tend=627.3976) ``` diff --git a/requirements.txt b/requirements.txt index 4fe9e43..800feb1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,8 +1,9 @@ fastprogress fire +# humecodec >=0.7.1 corrects Vorbis post-seek pts (audio_codec _seek_unreliable) +humecodec>=0.7.1 numpy polars>=1.36.1 pyarrow>=20 torch -torchaudio # torchcodec – optional, it causes serious performance regressions diff --git a/scripts/test_audio_backends.sh b/scripts/test_audio_backends.sh new file mode 100755 index 0000000..48bf02f --- /dev/null +++ b/scripts/test_audio_backends.sh @@ -0,0 +1,14 @@ +#!/usr/bin/env bash +# Run the full test suite in isolated environments, one per decoder backend. +# +# Usage: +# ./scripts/test_audio_backends.sh +# +# Each environment installs only one decoder backend so we verify that the +# interfaces work correctly regardless of which backend is present. + +TEST_COMMAND=${1:-"python -m tests"} + +parallel --tag --lb \ + "uv run --isolated --with {} $TEST_COMMAND" \ + ::: humecodec "torchaudio<2.9" torchcodec diff --git a/tests.py b/tests.py index 39a9acb..41ffc57 100644 --- a/tests.py +++ b/tests.py @@ -2,7 +2,7 @@ import unittest import wsds -from wsds import ws_dataset, ws_shard, ws_sink +from wsds import ws_dataset, ws_shard, ws_sink, ws_audio, audio_codec def load_tests(loader, tests, ignore): @@ -10,6 +10,8 @@ def load_tests(loader, tests, ignore): tests.addTests(doctest.DocTestSuite(ws_dataset)) tests.addTests(doctest.DocTestSuite(ws_shard)) tests.addTests(doctest.DocTestSuite(ws_sink)) + tests.addTests(doctest.DocTestSuite(ws_audio)) + tests.addTests(doctest.DocTestSuite(audio_codec)) tests.addTests(doctest.DocFileSuite("README.md")) return tests diff --git a/wsds/audio_codec.py b/wsds/audio_codec.py index 2c64c24..62eb9d4 100644 --- a/wsds/audio_codec.py +++ b/wsds/audio_codec.py @@ -2,236 +2,324 @@ This module contains all audio encoding/decoding logic, separated from the data model layer in ws_audio.py. It provides: -- Decoder backends (TorchFFmpegAudioDecoder, CompatAudioDecoder) -- A factory for creating decoders with automatic backend selection -- MP3 encoding with multi-backend fallback +- AudioDecoder: unified decoder with automatic backend selection (humecodec or torchaudio) +- encode_audio(): multi-backend encoder (humecodec -> torchcodec -> torchaudio) - HTML audio rendering utility """ from __future__ import annotations import io +import traceback import typing import pyarrow as pa -def to_filelike(src: typing.Any) -> typing.BinaryIO: - """Coerces files, byte-strings and PyArrow binary buffers into file-like objects.""" - if hasattr(src, "read"): # an open file - return src - # if not an open file then we assume some kind of binary data in memory - if hasattr(src, "as_buffer"): # PyArrow binary data - return pa.BufferReader(src.as_buffer()) - return io.BytesIO(src) +class AudioDecoder: + """Unified audio decoder that works with humecodec or torchaudio backends.""" + def __init__(self, reader, metadata, sample_rate, codec_delay=0): + self.reader = reader + self.metadata = metadata + self.sample_rate = sample_rate + self.debug = False + self.codec_delay = codec_delay + self.init_skip_samples = getattr(metadata, 'start_skip_samples', 0) or 0 + # Codecs where flush produces unreliable output (wrong skip_samples, + # wrong frame sizes). For these, always read from the start and trim. + # (vorbis was here too, but its post-seek transition-frame pts is now + # corrected in humecodec's StreamProcessor, so it seeks accurately.) + codec_name = getattr(metadata, 'codec', '') or '' + self._seek_unreliable = codec_name in ('wmav2', 'wmapro') + # Raw MPEG audio formats: timestamp seek does sequential scan, + # byte-offset seek with our own index is much faster. + self._use_byte_index = codec_name in ('mp3', 'mp2', 'mp1') + self._packet_index = None + + def _build_index(self): + """Build a sparse packet index for byte-offset seeking.""" + if self._packet_index is not None: + return + try: + idx = self.reader.build_packet_index( + self.reader.default_audio_stream, 128 * 1024) + if idx and len(idx) > 1: + self._packet_index = idx + except Exception: + self._packet_index = [] + + def _indexed_seek(self, target_time): + """Seek via byte offset using the packet index. Returns the index entry's PTS or None.""" + self._build_index() + if not self._packet_index: + return None + # Find last entry with pts <= target_time + best = self._packet_index[0] + for entry in self._packet_index: + if entry.pts_seconds <= target_time: + best = entry + else: + break + self.reader.seek_to_byte_offset(best.pos) + return best.pts_seconds + + def get_samples_played_in_range(self, tstart=0, tend=None, margin=.25): + import torch -class TorchFFmpegAudioDecoder: - def __init__(self, src, sample_rate): - from torchffmpeg import MediaDecoder + chunk = True + while chunk is not None: + (chunk,) = self.reader.pop_chunks() - if hasattr(src, "_optimal_read_size"): - buffer_size = src._optimal_read_size + # For short seeks and unreliable codecs, read from the start. + # This avoids seek accuracy issues for tstart < 5s (tiny cost) and + # codec flush bugs for wmav2/wmapro/vorbis. + read_from_start = self._seek_unreliable or tstart < 5.0 + + # Only adjust for start_skip_samples when actually seeking — when + # reading from start, the decoder applies skip_samples automatically. + seek_adj = 0.0 + index_pts = None + if not read_from_start: + # For raw MPEG formats, use indexed byte seek (fast, avoids sequential scan). + # No seek_adj needed: the index PTS and decoded audio are both in + # the raw timeline (skip_samples is not applied after byte seek). + if self._use_byte_index: + index_pts = self._indexed_seek(tstart - margin) + else: + # Timestamp seek: the demuxer applies start_skip_samples at + # pts=0 but not after seeking, so adjust tstart to compensate. + seek_adj = self.init_skip_samples / self.metadata.sample_rate + tstart += seek_adj + if tend is not None: + tend += seek_adj + + if index_pts is None: + # Fall back to timestamp seek (or read from start) + seek_target = 0.0 if read_from_start else max(0, tstart - margin) + self.reader.seek(seek_target, "key") + + chunks = [] + more_data = True + while more_data: + if self.reader.fill_buffer() == 1: + more_data = False + (chunk,) = self.reader.pop_chunks() + chunks.append(chunk) + if tend is not None: + chunk_end_pts = chunk.pts + chunk.shape[0] / self.sample_rate + if index_pts is not None: + # PTS not updated by demuxer after byte seek — estimate from index + elapsed = sum(c.shape[0] for c in chunks) / self.sample_rate + chunk_end_pts = index_pts + elapsed + if chunk_end_pts > tend + margin: + break + + # Determine the reference PTS for trimming + if read_from_start: + chunk0_pts = 0.0 + elif index_pts is not None: + # Byte seek: demuxer PTS is stale, use our index entry + chunk0_pts = index_pts else: - buffer_size = 128 * 1024 - self.src = src - self.reader = MediaDecoder(to_filelike(self.src), buffer_size=buffer_size) - self.metadata = self.reader.get_src_stream_info(self.reader.default_audio_stream) + chunk0_pts = chunks[0].pts + prefix = round(tstart * self.sample_rate) - round(chunk0_pts * self.sample_rate) + + if self.debug: + import torch as _t + total_samples = sum(c.shape[0] for c in chunks) + print(f" [decode] codec={self.metadata.codec} sr={self.sample_rate} " + f"tstart_orig={tstart - seek_adj:.4f} tstart_adj={tstart:.4f} " + f"seek_adj={seek_adj:.6f} (init_skip={self.init_skip_samples} codec_delay={self.codec_delay}) " + f"chunk0.pts={chunks[0].pts:.6f} chunk0_pts_used={chunk0_pts:.6f} " + f"n_chunks={len(chunks)} total_samples={total_samples} prefix={prefix}", flush=True) - if sample_rate is None: - sample_rate = int(self.metadata.sample_rate) + if prefix < 0: + if self.debug: + print(f" [trim] negative prefix {prefix}, clamping to 0", flush=True) + prefix = 0 + # Unwrap humecodec Chunk (a torch.Tensor subclass) to its plain `_elem` + # tensor before cat: otherwise every torch op on a Chunk goes through + # __torch_dispatch__ -> tree_map(unwrap, ...), which allocates ~14 cyclic + # pytree objects per decode and feeds the GC-collection latency spikes. + samples = torch.cat([getattr(c, "_elem", c) for c in chunks]) + if tend is not None: + return samples[prefix : prefix + round(tend * self.sample_rate) - round(tstart * self.sample_rate)].mT + else: + return samples[prefix:].mT - self.sample_rate = sample_rate + def add_seek_points(self, positions, pts_seconds): + """Seed the demuxer's seek index with precomputed (byte position, pts) + pairs so a subsequent timestamp seek() brackets the target and converges + in ~1 read instead of a full binary/secant search across the file + (dramatic for Ogg/Vorbis, which has no container index). `positions` are + byte offsets in the input (blob-relative when decoding via a LazyBuffer + over the audio blob); `pts_seconds` are in seconds. Accuracy is + unchanged — the seek still reads the landing page to position exactly. - self.reader.add_basic_audio_stream( - frames_per_chunk=int(32 * sample_rate), - sample_rate=sample_rate, - decoder_option={"threads": "4", "thread_type": "frame"}, - ) + Requires humecodec >= 0.8 (the `add_seek_points` backend method); returns + False and no-ops on older builds or the torchcodec backend. + """ + fn = getattr(self.reader, "add_seek_points", None) + if fn is None: + return False + fn([int(p) for p in positions], [float(t) for t in pts_seconds]) + return True - def get_samples_played_in_range(self, tstart=0, tend=None): - import torch - self.reader.seek(max(0, tstart - 1), "key") - - if tend is None: - chunks = [] - more_data = True - while more_data: - if self.reader.fill_buffer() == 1: - more_data = False - (chunk,) = self.reader.pop_chunks() - if chunk is not None: - chunks.append(chunk) - prefix = int((tstart - chunks[0].pts) * self.sample_rate) - if prefix < 0: - prefix = 0 - return torch.cat(chunks)[prefix:].mT - - self.reader.fill_buffer() - (chunk,) = self.reader.pop_chunks() - prefix = int((tstart - chunk.pts) * self.sample_rate) - if prefix < 0: - prefix = 0 - if tend: - samples = chunk[prefix : prefix + int((tend - tstart) * self.sample_rate)].mT - else: - samples = chunk[prefix:].mT - while chunk is not None: - (chunk,) = self.reader.pop_chunks() - return samples +def _create_reader_humecodec(src, buffer_size): + from humecodec import MediaDecoder + reader = MediaDecoder(src=src, buffer_size=buffer_size) + metadata = reader.get_src_stream_info(reader.default_audio_stream) + return reader, metadata -class CompatAudioDecoder: - def __init__(self, src, sample_rate): - import torchaudio - if not hasattr(torchaudio, "io"): - raise ImportError("You need either torchaudio<2.9 or torchcodec installed") - self.src = src - if hasattr(src, "_optimal_read_size"): - buffer_size = src._optimal_read_size - else: - buffer_size = 128 * 1024 - self.reader = torchaudio.io.StreamReader(src=to_filelike(self.src), buffer_size=buffer_size) - self.metadata = self.reader.get_src_stream_info(0) +def _create_reader_torchaudio(src, buffer_size): + from torchaudio.io import StreamReader - if sample_rate is None: - sample_rate = self.metadata.sample_rate + reader = StreamReader(src=src, buffer_size=buffer_size) + metadata = reader.get_src_stream_info(reader.default_audio_stream) + return reader, metadata - self.sample_rate = sample_rate - # fetch 32 seconds because we likely need 30s at maximum but the seeking may be imprecise (and we seek 1s early) - # FIXME: check if we can get away with some better settings here (-1, maybe 10s + concatenate the chunks in a loop) - self.reader.add_basic_audio_stream( - frames_per_chunk=int(32 * sample_rate), - sample_rate=sample_rate, - decoder_option={"threads": "4", "thread_type": "frame"}, - ) - - def get_samples_played_in_range(self, tstart=0, tend=None): - # rought seek - self.reader.seek(max(0, tstart - 1), "key") - - if tend is None: - import torch - - chunks = [] - more_data = True - while more_data: - if self.reader.fill_buffer() == 1: - more_data = False - (chunk,) = self.reader.pop_chunks() - chunks.append(chunk) - prefix = int((tstart - chunks[0].pts) * self.sample_rate) - if prefix < 0: - prefix = 0 - return torch.cat(chunks)[prefix:].mT - - self.reader.fill_buffer() - (chunk,) = self.reader.pop_chunks() - # tight crop (seems accurate down to 1 sample in my tests) - prefix = int((tstart - chunk.pts) * self.sample_rate) - if prefix < 0: - prefix = 0 - if tend: - samples = chunk[prefix : prefix + int((tend - tstart) * self.sample_rate)].mT - else: - samples = chunk[prefix:].mT - # clear out any remaining data - while chunk is not None: - (chunk,) = self.reader.pop_chunks() - return samples +def _create_decoder_torchcodec(src, sample_rate): + """Create a torchcodec-backed decoder that matches the AudioDecoder interface.""" + from types import SimpleNamespace + from torchcodec.decoders import AudioDecoder as TorchcodecDecoder -def create_decoder(src, sample_rate=None): - """Factory: tries torchffmpeg -> torchcodec -> torchaudio, returns a decoder instance. + # torchcodec accepts bytes but not BytesIO + decoder = TorchcodecDecoder(src, sample_rate=sample_rate) + metadata = decoder.metadata - Args: - src: A file-like object or bytes-like source for audio data. - sample_rate: Optional target sample rate for resampling. + class TorchcodecAdapter: + def __init__(self): + self.metadata = metadata + self.sample_rate = sample_rate if sample_rate is not None else int(metadata.sample_rate) - Returns: - A decoder instance with .metadata, .sample_rate, and .get_samples_played_in_range() interface. - """ - try: - from torchffmpeg import MediaDecoder as _ # noqa: F401 + def get_samples_played_in_range(self, tstart=0, tend=None): + return decoder.get_samples_played_in_range(tstart, tend) - AudioDecoder = TorchFFmpegAudioDecoder - except ImportError: - try: - from torchcodec.decoders import AudioDecoder - except ImportError: - AudioDecoder = CompatAudioDecoder + return TorchcodecAdapter() - return AudioDecoder(src, sample_rate=sample_rate) +_STREAMING_BACKENDS = [ + (_create_reader_humecodec, "humecodec"), + (_create_reader_torchaudio, "torchaudio.io"), +] -def decode_segment(src, start=0, end=None, sample_rate=None): - """One-shot decode: creates decoder, reads segment, returns tensor with .sample_rate attr. +_chosen_backend = None - Handles MP3 skip_samples compensation automatically. + +def create_decoder(src, sample_rate=None): + """Factory: tries humecodec -> torchaudio -> torchcodec, returns a decoder instance. Args: - src: Audio source (file-like, bytes, or PyArrow buffer). - start: Start time in seconds. - end: End time in seconds (None for rest of file). - sample_rate: Optional target sample rate. + src: A file-like object for audio data. + sample_rate: Optional target sample rate for resampling. Returns: - A torch.Tensor with a .sample_rate attribute. + A decoder with .metadata, .sample_rate, and .get_samples_played_in_range(). """ - filelike = to_filelike(src) - decoder = create_decoder(filelike, sample_rate) - - skip_samples = 0 - if decoder.metadata.codec == "mp3": - skip_samples = 1105 + global _chosen_backend + + buffer_size = getattr(src, "_optimal_read_size", 128 * 1024) + + if _chosen_backend is not None: + if _chosen_backend == "torchcodec": + return _create_decoder_torchcodec(src, sample_rate) + reader, metadata = _chosen_backend(src, buffer_size) + else: + for factory, module in _STREAMING_BACKENDS: + try: + reader, metadata = factory(src, buffer_size) + _chosen_backend = factory + break + except ImportError: + continue + else: + # Fall back to torchcodec (different API, no streaming reader) + try: + decoder = _create_decoder_torchcodec(src, sample_rate) + _chosen_backend = "torchcodec" + return decoder + except ImportError: + raise ImportError("Neither humecodec, torchaudio, nor torchcodec is installed.") if sample_rate is None: - sample_rate = decoder.metadata.sample_rate + sample_rate = int(metadata.sample_rate) - seek_adjustment = skip_samples / sample_rate if start > 0 else 0 - samples = decoder.get_samples_played_in_range( - start + seek_adjustment, end + seek_adjustment if end is not None else None + reader.add_basic_audio_stream( + frames_per_chunk=int(1 * sample_rate), + sample_rate=sample_rate, + decoder_option={"threads": "4", "thread_type": "frame"}, ) - if hasattr(samples, "data"): - samples = samples.data - samples.sample_rate = sample_rate - return samples + + # Get codec_delay from the decoder (available after add_audio_stream opens the codec) + codec_delay = 0 + try: + out_info = reader.get_out_stream_info(0) + codec_delay = getattr(out_info, 'codec_delay', 0) or 0 + except Exception: + pass + + return AudioDecoder(reader, metadata, sample_rate, codec_delay=codec_delay) + -def encode_mp3(samples) -> bytes: - """Encode a torch tensor to MP3 bytes. +def encode_audio(samples, format="mp3", sample_rate=None, bitrate=None) -> bytes: + """Encode a torch tensor to audio bytes. - Tries torchffmpeg -> torchcodec -> torchaudio as encoder backends. + Tries humecodec -> torchcodec -> torchaudio as encoder backends. + + >>> from wsds import WSDataset + >>> audio = WSDataset("librilight/source")[0].get_audio() + >>> samples = audio.read_segment(start=0, end=2.0, sample_rate=16000) + >>> mp3 = encode_audio(samples, format="mp3") + >>> mp3[:3] == b"ID3" or mp3[:2] in (b"\\xff\\xfb", b"\\xff\\xf3") + True + >>> ogg = encode_audio(samples, format="ogg") # doctest: +SKIP + >>> ogg[:4] == b"OggS" # doctest: +SKIP + True Args: samples: A torch.Tensor with a .sample_rate attribute. Shape: (channels, frames). + format: Output format, e.g. "mp3", "ogg" (Opus). Default: "mp3". + sample_rate: Target sample rate (defaults to samples.sample_rate). + bitrate: Bitrate in bps. Only used for formats that support it (e.g. Opus). Returns: - MP3-encoded bytes. + Encoded audio bytes. """ + if sample_rate is None: + sample_rate = int(samples.sample_rate) + out = io.BytesIO() try: - from torchffmpeg import MediaEncoder + from humecodec import MediaEncoder - sample_rate = int(samples.sample_rate) - # samples is (channels, frames), write_audio_chunk expects (frames, channels) waveform = samples.mT.float().contiguous() - enc = MediaEncoder(out, "mp3") - enc.add_audio_stream(sample_rate=sample_rate, num_channels=waveform.size(1), format="flt") + enc = MediaEncoder(out, format) + stream_kwargs = dict(sample_rate=sample_rate, num_channels=waveform.size(1), format="flt") + if format == "ogg": + from humecodec import CodecConfig + + stream_kwargs.update(encoder="libopus", encoder_format="flt") + if bitrate: + stream_kwargs["codec_config"] = CodecConfig(bit_rate=bitrate) + enc.add_audio_stream(**stream_kwargs) with enc.open(): enc.write_audio_chunk(0, waveform) except ImportError: try: from torchcodec.encoders import AudioEncoder - AudioEncoder(samples, sample_rate=int(samples.sample_rate)).to_file_like(out, "mp3") + AudioEncoder(samples, sample_rate=sample_rate).to_file_like(out, format) except ImportError: import torchaudio - torchaudio.save(out, samples, int(samples.sample_rate), format="mp3") + torchaudio.save(out, samples, sample_rate, format=format) return out.getvalue() @@ -247,5 +335,5 @@ def audio_to_html(samples) -> str: """ import base64 - mp3_data = base64.b64encode(encode_mp3(samples)).decode("ascii") + mp3_data = base64.b64encode(encode_audio(samples, format="mp3")).decode("ascii") return f'' diff --git a/wsds/ws_audio.py b/wsds/ws_audio.py index 1dd7351..6b22a2f 100644 --- a/wsds/ws_audio.py +++ b/wsds/ws_audio.py @@ -3,31 +3,32 @@ import typing from dataclasses import dataclass -from .audio_codec import audio_to_html, create_decoder, encode_mp3, to_filelike +from .audio_codec import audio_to_html, create_decoder, encode_audio +from .pupyarrow import pupyarrow -def load_segment(src, start, end, sample_rate=None): - """Efficiently loads an audio segment from `src` (see below) `tstart` to `tend` seconds while - optionally resampling it to `sample_rate`. - - `src` can be one of: - - a file-like object - - a byte string - - a PyArrow binary buffer in memory""" - return AudioReader(src).read_segment(start, end, sample_rate=sample_rate) - @dataclass() -class AudioReader: - """A lazy seeking-capable audio reader for random-access to recordings stored in wsds shards.""" +class WSAudioEpisode: + """A lazy seeking-capable audio reader for random-access to recordings stored in wsds shards. + + >>> from wsds import WSDataset + >>> ds = WSDataset("librilight/source") + >>> audio = ds[0].get_audio() + >>> audio.load().shape + torch.Size([1, 17884909]) + >>> audio.read_segment(start=2, end=5).shape + torch.Size([1, 48000]) + >>> audio.read_segment(start=2, end=5, sample_rate=8000).shape + torch.Size([1, 24000]) + """ src: typing.Any _decoder: typing.Any = None _sample_rate: int | None = None - skip_samples: int = 0 - + _seek_index: typing.Any = None # (positions_blob_relative, pts_seconds) or None def __repr__(self): - return f"AudioReader(src={type(self.src)}, sample_rate={self._sample_rate})" + return f"WSAudioEpisode(src={type(self.src)}, sample_rate={self._sample_rate})" def unwrap(self): """Return the raw audio bytes""" @@ -35,27 +36,45 @@ def unwrap(self): return self.src.as_buffer().to_pybytes() elif isinstance(self.src, (bytes, bytearray)): return self.src + elif isinstance(self.src, pupyarrow.LazyBuffer): + return self.src.read() else: - raise TypeError(f"Unsupported AudioReader src type: {type(self.src)}") + raise TypeError(f"Unsupported src type: {type(self.src)}") + + to_bytes = unwrap + + def set_seek_index(self, positions, pts_seconds): + """Attach a precomputed seek index so seeks skip the expensive on-open + scan: `build_packet_index` for mp3/mp2/mp1 (which reads the WHOLE episode + — ~200-400ms for a long spotify mp3, the dominant deep-seek cost) and the + Ogg/Vorbis demuxer bisection. `positions` are byte offsets in the audio + blob (blob-relative); `pts_seconds` the matching timestamps. Applied on + the next decoder (re)creation.""" + self._seek_index = ([int(p) for p in positions], [float(t) for t in pts_seconds]) + if self._decoder is not None: + self._apply_seek_index() + + def _apply_seek_index(self): + if not self._seek_index or self._decoder is None: + return + pos, pts = self._seek_index + d = self._decoder + if getattr(d, "_use_byte_index", False): # mp3/mp2/mp1 byte-index path + from types import SimpleNamespace + d._packet_index = [SimpleNamespace(pts_seconds=t, pos=p) for p, t in zip(pos, pts)] + else: # vorbis/etc: seed the demuxer index (humecodec>=0.8) + fn = getattr(d, "add_seek_points", None) + if fn is not None: + fn(pos, pts) def get_decoder(self, sample_rate=None): """Lazily creates/caches decoder via audio_codec.create_decoder().""" - sample_rate_switch = False - if self._sample_rate is not None: - sample_rate_switch = self._sample_rate != sample_rate - - if self._decoder is None or sample_rate_switch: - decoder = create_decoder(to_filelike(self.src), sample_rate=sample_rate) - # mp3 has encoder delays that are not handled well when seeking - if decoder.metadata.codec == "mp3": - self.skip_samples = 1105 - - if sample_rate is None: - sample_rate = decoder.metadata.sample_rate - - self._decoder = decoder - self._sample_rate = sample_rate - + requested_sr = sample_rate or (self._decoder and self._decoder.metadata.sample_rate) + if self._decoder is None or requested_sr != self._sample_rate: + self.src.seek(0) + self._decoder = create_decoder(self.src, sample_rate=sample_rate) + self._sample_rate = sample_rate or self._decoder.metadata.sample_rate + self._apply_seek_index() return self._decoder, self._sample_rate @property @@ -70,14 +89,9 @@ def sample_rate(self): def read_segment(self, start=0, end=None, sample_rate=None): decoder, sample_rate = self.get_decoder(sample_rate) - seek_adjustment = self.skip_samples / sample_rate if start > 0 else 0 - _samples = decoder.get_samples_played_in_range( - start + seek_adjustment, end + seek_adjustment if end is not None else None - ) - if hasattr(_samples, "data"): - samples = _samples.data - else: - samples = _samples + samples = decoder.get_samples_played_in_range(start, end) + if hasattr(samples, "data"): + samples = samples.data samples.sample_rate = sample_rate return samples @@ -91,56 +105,60 @@ def _repr_html_(self): def _display_(self): import marimo - return marimo.audio(encode_mp3(self.read_segment())) + return marimo.audio(encode_audio(self.read_segment())) @dataclass(frozen=True) -class WSAudio: - """A lazy reference to a single sample from a segmented audio file.""" +class WSAudioSegment: + """A lazy reference to a single sample from a segmented audio file. + """ - audio_reader: AudioReader + episode: WSAudioEpisode tstart: float tend: float + def __repr__(self) -> str: + return f"WSAudioSegment(episode={self.episode}, tstart={self.tstart!s}, tend={self.tend!s})" + @property def duration(self) -> float: """Duration of the audio segment in seconds.""" return self.tend - self.tstart - def with_context(self, before: float = 0, after: float = 0) -> "WSAudio": - """Return a new WSAudio with expanded timestamps to include surrounding context. + def with_context(self, before: float = 0, after: float = 0) -> "WSAudioSegment": + """Return a new WSAudioSegment with expanded timestamps to include surrounding context. Args: before: Seconds of context to add before the segment start (will not go below 0) after: Seconds of context to add after the segment end Returns: - A new WSAudio instance with adjusted timestamps + A new WSAudioSegment instance with adjusted timestamps """ - return WSAudio( - audio_reader=self.audio_reader, + return WSAudioSegment( + episode=self.episode, tstart=max(0, self.tstart - before), tend=self.tend + after, ) - def with_timestamps(self, tstart: float | None = None, tend: float | None = None) -> "WSAudio": - """Return a new WSAudio with modified timestamps. + def with_timestamps(self, tstart: float | None = None, tend: float | None = None) -> "WSAudioSegment": + """Return a new WSAudioSegment with modified timestamps. Args: tstart: New start time in seconds (None to keep current) tend: New end time in seconds (None to keep current) Returns: - A new WSAudio instance with the specified timestamps + A new WSAudioSegment instance with the specified timestamps """ - return WSAudio( - audio_reader=self.audio_reader, + return WSAudioSegment( + episode=self.episode, tstart=tstart if tstart is not None else self.tstart, tend=tend if tend is not None else self.tend, ) def load(self, sample_rate=None, pad_to_seconds=None): - samples = self.audio_reader.read_segment(self.tstart, self.tend, sample_rate) + samples = self.episode.read_segment(self.tstart, self.tend, sample_rate) sample_rate = samples.sample_rate if pad_to_seconds is not None: import torch @@ -152,7 +170,7 @@ def load(self, sample_rate=None, pad_to_seconds=None): @property def metadata(self): - return self.audio_reader.metadata + return self.episode.metadata def _repr_html_(self): return audio_to_html(self.load()) @@ -160,4 +178,4 @@ def _repr_html_(self): def _display_(self): import marimo - return marimo.audio(encode_mp3(self.load())) + return marimo.audio(encode_audio(self.load())) diff --git a/wsds/ws_audio_index.py b/wsds/ws_audio_index.py new file mode 100644 index 0000000..eb50a24 --- /dev/null +++ b/wsds/ws_audio_index.py @@ -0,0 +1,311 @@ +"""Audio seek index generation and efficient segment loading using humecodec. + +This module provides: +- `generate_audio_seek_index`: Scans audio shards using pupyarrow to get byte + offsets, then uses humecodec's `build_packet_index` to build a seek index + for each audio blob. The index is stored in a new wsds shard. +- `load_audio_segment`: Given a WSSample (containing seek index columns) and a + timestamp + duration, opens the shard file directly via `shard.get_reader()` + and decodes only the needed portion using `process_packet()`. + +The seek index shard contains these columns per sample: +- `__key__`: The sample key (matching the audio shard) +- `seek_index_audio_offset`: int64 byte offset of the audio blob inside the audio shard +- `seek_index_audio_length`: int64 byte length of the audio blob +- `seek_index_positions`: list[int64] absolute byte positions inside the shard, + spaced every ~512kB +- `seek_index_pts_seconds`: list[float64] corresponding timestamps in seconds +- `seek_index_duration`: float64 total audio duration in seconds +""" + +from __future__ import annotations + +import io +import typing +from pathlib import Path + +import humecodec +import polars as pl +import torch + +from .pupyarrow.file_reader import LocalFileReader +from .pupyarrow.pupyarrow import FeatherFile, LazyBinaryArray, LazyBuffer + +SEEK_RESOLUTION_BYTES = 512 * 1024 # 512kB between seek points + +# Header/footer bytes cached in the index and served locally so a decoder's +# open fetches nothing from the store — a cold seek then touches only the target +# blocks. The decoder is opened with a SMALL avio buffer (DECODE_BUFFER_BYTES): +# find_stream_info then reads only ~8kB of header (a big buffer inflates it to +# ~132kB), and no read-ahead is needed because the block-cache layer already +# provides it (whole 128kB blocks, served locally). The Ogg duration probe still +# scans back a ~64kB window (FOOTER_PROBE_BYTES) but only needs the real last +# page, so we cache the last FOOTER_CACHE_BYTES and zero-fill the rest. +HEADER_CACHE_BYTES = 8192 +FOOTER_CACHE_BYTES = 8192 +FOOTER_PROBE_BYTES = 65536 +DECODE_BUFFER_BYTES = 4096 + + +def _measure_header_len(audio_bytes, buffer_size=DECODE_BUFFER_BYTES, cap=131072): + """Bytes the decoder actually reads at the head during open (find_stream_info + + first-packet confirm) with `buffer_size` — cache exactly this so a cold + open touches 0 shard blocks. Self-adjusts to setup-header size (codebook + complexity) vs a fixed guess. The consumer opens with the same buffer_size. + Falls back to HEADER_CACHE_BYTES on failure.""" + reads = [] + + class _T: + def __init__(s): s.pos = 0; s.n = len(audio_bytes) + def read(s, k=-1): + if k is None or k < 0: k = s.n - s.pos + k = min(k, s.n - s.pos) + d = audio_bytes[s.pos:s.pos + k]; reads.append((s.pos, len(d))); s.pos += len(d); return d + read1 = read + def seek(s, o, w=0): s.pos = o if w == 0 else (s.pos + o if w == 1 else s.n + o); return s.pos + def tell(s): return s.pos + def size(s): return s.n + def readable(s): return True + def seekable(s): return True + + try: + r = humecodec.MediaDecoder(src=_T(), buffer_size=buffer_size) + info = r.get_src_stream_info(r.default_audio_stream) + r.add_basic_audio_stream(frames_per_chunk=int(info.sample_rate), + sample_rate=int(info.sample_rate)) + except Exception: + return min(cap, len(audio_bytes), HEADER_CACHE_BYTES) + half = max(1, len(audio_bytes) // 2) + return min(cap, max((a + n for a, n in reads if a < half), default=HEADER_CACHE_BYTES)) + + +def _ogg_footer(buf, search_cap=65536, fallback=FOOTER_CACHE_BYTES): + """The real footer the Ogg duration probe needs: bytes from the last page + carrying a VALID granule (>=0) to EOF (typically ~2-4kB). Non-Ogg / not + found -> last `fallback` bytes. Searches only the last `search_cap` bytes.""" + if buf[:4] != b"OggS": + return bytes(buf[-min(fallback, len(buf)):]) + end = len(buf) + lo = max(0, end - search_cap) + pos = buf.rfind(b"OggS", lo) + while pos != -1: + # granule position = int64 LE at [pos+6, pos+14); -1 (0xFFFF..) == no packet + if pos + 14 <= end and int.from_bytes(buf[pos + 6:pos + 14], "little", signed=True) >= 0: + return bytes(buf[pos:]) + pos = buf.rfind(b"OggS", lo, pos) + return bytes(buf[-min(fallback, len(buf)):]) + + +def generate_audio_seek_index( + audio_shard_path: str | Path, + output_path: str | Path, + resolution: int = SEEK_RESOLUTION_BYTES, + key_column: str = "__key__", + audio_column: str = "audio", + compression: str | None = "zstd", +): + """Generate a seek index shard for an audio shard. + + For each audio sample in the shard, opens it with humecodec.MediaDecoder + and calls `build_packet_index(resolution=...)` to get a sparse index of + byte positions and pts values. These are stored as absolute shard-file + offsets so that a reader can later seek directly within the shard. + + Args: + audio_shard_path: Path to the source .wsds shard containing audio. + output_path: Path for the output seek index .wsds shard. + resolution: Minimum byte distance between index entries (default 512KB). + key_column: Name of the key column in the source shard. + audio_column: Name of the audio column in the source shard. + compression: Compression for the output shard (default "zstd"). + """ + from .ws_sink import WSSink + + audio_shard_path = Path(audio_shard_path) + output_path = Path(output_path) + + reader = LocalFileReader(audio_shard_path) + feather = FeatherFile(reader) + + rows = [] + + for batch_idx in range(feather.num_record_batches): + batch = feather.record_batch(batch_idx) + key_col = batch.column(key_column) + audio_col = batch.column(audio_column) + + if not isinstance(audio_col, LazyBinaryArray): + raise TypeError(f"Expected binary column for '{audio_column}', got {type(audio_col).__name__}") + + # Absolute file offset of the data buffer backing this column + data_buf_offset = audio_col._data_buffer._offset + + for i in range(batch.num_rows): + key = key_col[i] + if isinstance(key, bytes): + key = key.decode("utf-8") + + # Absolute byte offset and length of this audio element in the shard file + elem_start = int(audio_col.offsets[i]) + elem_end = int(audio_col.offsets[i + 1]) + audio_offset = data_buf_offset + elem_start + audio_length = elem_end - elem_start + + # Read the audio blob and build a packet index via humecodec + audio_buf = audio_col[i] + audio_bytes = audio_buf._read_all() + + decoder = humecodec.MediaDecoder(io.BytesIO(audio_bytes), buffer_size=len(audio_bytes)) + decoder.add_audio_stream(frames_per_chunk=-1) + packet_index = decoder.build_packet_index(resolution=resolution) + + # Store positions as absolute shard-file offsets, pts as seconds + seek_positions = [audio_offset + entry.pos for entry in packet_index] + seek_pts_seconds = [entry.pts_seconds for entry in packet_index] + + # Compute total duration from the last packet + if packet_index: + last = packet_index[-1] + duration = last.pts_seconds + last.duration_seconds + else: + duration = 0.0 + + rows.append( + { + key_column: key, + "seek_index_audio_offset": audio_offset, + "seek_index_audio_length": audio_length, + "seek_index_positions": seek_positions, + "seek_index_pts_seconds": seek_pts_seconds, + "seek_index_duration": duration, + # served locally at decode time so open touches 0 shard blocks + "seek_index_header": audio_bytes[:_measure_header_len(audio_bytes)], + "seek_index_footer": _ogg_footer(audio_bytes), # real last page (~2-4kB) + } + ) + + feather.close() + + with WSSink(str(output_path), compression=compression) as sink: + for row in rows: + sink.write(row) + + +class _HFOverlayReader: + """Wraps a FileReader and serves the audio blob's first ``len(header)`` and + last ``len(footer)`` bytes from in-memory buffers cached in the seek index, + so a decoder's open-time reads (header/setup + Ogg duration probe) never hit + the underlying store. Offsets are absolute; the blob spans + ``[audio_offset, audio_offset + audio_length)``. Non-read attributes are + delegated to the wrapped reader.""" + + def __init__(self, reader, audio_offset, audio_length, header, footer): + self._reader = reader + self._ao = int(audio_offset) + self._al = int(audio_length) + self._header = header + self._footer = footer + self._H = len(header) + self._F = len(footer) + + def read(self, offset: int, length: int) -> bytes: + s = offset - self._ao # blob-relative start + e = s + length + L = self._al + fc_start = L - self._F # real footer bytes + # zeros cover the duration-probe window before the real last page + fz_start = max(self._H, L - FOOTER_PROBE_BYTES) if self._F else L + if s < 0 or length <= 0 or (s >= self._H and e <= fz_start): + return self._reader.read(offset, length) # pure middle / OOB + if e <= self._H: + return self._header[s:e] # pure header + if s >= fc_start: + return self._footer[s - fc_start:e - fc_start] # pure footer + out = bytearray(length) # spans regions + he = min(e, self._H) + if s < he: + out[0:he - s] = self._header[s:he] + fs = max(s, fc_start) + if self._F and fs < e: + out[fs - s:length] = self._footer[fs - fc_start:e - fc_start] + # [fz_start, fc_start) left as zeros (duration probe, never fetched) + ms, me = max(s, self._H), min(e, fz_start) + if ms < me: + out[ms - s:me - s] = self._reader.read(self._ao + ms, me - ms) + return bytes(out) + + def __getattr__(self, name): + return getattr(self._reader, name) + + +def load_audio_segment( + sample: typing.Any, + timestamp: float, + duration: float, + sample_rate: int | None = None, + audio_column: str = "audio", +) -> torch.Tensor: + """Efficiently load an audio segment using the seek index. + + Opens the shard file directly via `shard.get_reader()` and creates a + `LazyBuffer` view over just the audio blob region. Only the bytes needed + for header probing and the requested segment are read — no full-blob + download. Works for all shard sources (local, S3, Modal). + + Args: + sample: A WSSample (or dict-like) containing at minimum: + - "seek_index_audio_offset": int64 byte offset of the blob in the shard + - "seek_index_audio_length": int64 byte length of the blob + - "seek_index_positions": list[int64] absolute byte positions in the shard + - "seek_index_pts_seconds": list[float64] timestamps in seconds + timestamp: Start time in seconds. + duration: Duration in seconds. + sample_rate: Target sample rate for resampling. None keeps the native rate. + audio_column: Name of the audio column (used to locate the shard). + + Returns: + Torch tensor of shape (channels, samples) with decoded audio. + """ + from .audio_codec import create_decoder + + audio_offset = sample["seek_index_audio_offset"] + audio_length = sample["seek_index_audio_length"] + seek_positions = sample["seek_index_positions"] # absolute shard offsets + seek_pts_seconds = sample["seek_index_pts_seconds"] + + # LazyBuffer view over just the audio blob (blob-relative offsets); reads are + # sparse — only the header + the seeked segment are fetched. + column_dir, _ = sample.dataset.fields[audio_column][0] + shard = sample.dataset.get_shard(column_dir, sample.shard_ref) + file_reader = shard.get_reader() + + # Serve cached header/footer locally (if present) so the decoder's open-time + # reads fetch 0 blocks from the store — cold seek touches only target blocks. + def _hf(name): + try: + v = sample[name] + except (KeyError, TypeError): + return b"" + return bytes(v) if v else b"" + header, footer = _hf("seek_index_header"), _hf("seek_index_footer") + if header or footer: + file_reader = _HFOverlayReader(file_reader, audio_offset, audio_length, header, footer) + + audio_view = LazyBuffer(file_reader, audio_offset, audio_length) + # Small avio buffer: keeps find_stream_info's header read to ~8kB (a large + # buffer inflates it) — read-ahead is redundant with the block-cache layer. + audio_view._optimal_read_size = DECODE_BUFFER_BYTES + + # Per-codec decoder (mp4->moov/timestamp, vorbis->corrected granule seek, + # wma->read-from-start, mp3->byte index). Seed the demuxer index so the + # timestamp seek converges in ~1 read instead of scanning the whole file. + dec = create_decoder(audio_view, sample_rate) + rel_pos = [p - audio_offset for p in seek_positions] # blob-relative + dec.add_seek_points(rel_pos, seek_pts_seconds) + # Also seed the byte-index (mp3/mp2/mp1) path so it doesn't rescan to build. + if getattr(dec, "_use_byte_index", False): + from types import SimpleNamespace + dec._packet_index = [SimpleNamespace(pts_seconds=float(t), pos=int(p)) + for p, t in zip(rel_pos, seek_pts_seconds)] + + return dec.get_samples_played_in_range(timestamp, timestamp + duration) diff --git a/wsds/ws_dataset.py b/wsds/ws_dataset.py index c8fd72e..cf96590 100644 --- a/wsds/ws_dataset.py +++ b/wsds/ws_dataset.py @@ -39,8 +39,8 @@ class WSDataset: >>> sample = dataset["large/5304/the_tinted_venus_1408_librivox_64kb_mp3/tintedvenus_05_anstey_64kb_090"] >>> print(repr(sample["transcription_wslang_raw.txt"])) ' I will accompany you," she said.' - >>> sample['audio'] - WSAudio(audio_reader=AudioReader(src=, sample_rate=None), tstart=1040.2133, tend=1042.8413) + >>> sample['audio'].load().shape + torch.Size([1, 42049]) """ dataset_root: Path diff --git a/wsds/ws_decode.py b/wsds/ws_decode.py index 7d89979..e81a6f0 100644 --- a/wsds/ws_decode.py +++ b/wsds/ws_decode.py @@ -5,7 +5,7 @@ import numpy as np import pyarrow as pa -from .ws_audio import AudioReader +from .ws_audio import WSAudioEpisode AUDIO_FILE_KEYS = frozenset( [ @@ -32,6 +32,32 @@ def to_filelike(src: typing.Any) -> typing.BinaryIO: return io.BytesIO(src) +# --- ".arr": variable-length numeric arrays stored NATIVELY in pyarrow -------- +# Unlike ".npy" (a binary blob decoded with np.load, which parses the header via +# ast.literal_eval and allocates cyclic garbage every read), a ".arr" column is a +# native `list>` (2-D) or `list` (1-D). Reads are a +# zero-copy `flatten().to_numpy()` — ~10x faster, no per-read cyclic allocation. + +def arr_column_type(sample) -> pa.DataType: + """Native pyarrow type for a ".arr" column value (numpy 1-D or 2-D).""" + a = np.asarray(sample) + et = pa.from_numpy_dtype(a.dtype) + if a.ndim == 1: + return pa.list_(et) # list + if a.ndim == 2: + return pa.list_(pa.list_(et, a.shape[1])) # list> + raise ValueError(f".arr supports 1-D/2-D arrays, got ndim={a.ndim}") + + +def decode_arr(scalar, arrow_type) -> np.ndarray: + """Read a native ".arr" list scalar as a numpy array (no np.load / as_py).""" + inner = arrow_type.value_type + if pa.types.is_fixed_size_list(inner): + vals = scalar.values.flatten().to_numpy(zero_copy_only=False) + return vals.reshape(-1, inner.list_size) + return scalar.values.to_numpy(zero_copy_only=False) + + def decode_sample(column: str, data): """Decode a binary column value from a file-like object based on column name. @@ -50,7 +76,7 @@ def decode_sample(column: str, data): import json return json.load(fd) elif ext in AUDIO_FILE_KEYS: - return AudioReader(fd) + return WSAudioEpisode(fd) else: return fd.read() @@ -64,6 +90,9 @@ def encode_value(column: str, value): buf = io.BytesIO() np.save(buf, value) return buf.getvalue() + elif ext == "arr": + # keep the numpy array; WSSink builds the native list column (arr_column_type) + return np.asarray(value) elif ext == "pyd": return pickle.dumps(value) elif ext == "json": @@ -90,7 +119,7 @@ def get_audio(sample, audio_columns=None): audio_columns: Optional list of column names to try. Defaults to AUDIO_FILE_KEYS. Returns: - The audio value (typically an AudioReader or WSAudio). + The audio value (typically a WSAudioEpisode or WSAudioSegment). Raises: KeyError: If no audio column is found in the sample. diff --git a/wsds/ws_shard.py b/wsds/ws_shard.py index 39639cc..4d99b6c 100644 --- a/wsds/ws_shard.py +++ b/wsds/ws_shard.py @@ -7,8 +7,8 @@ import pyarrow as pa from .utils import WSShardMissingError -from .ws_audio import AudioReader, WSAudio -from .ws_decode import decode_sample +from .ws_audio import WSAudioEpisode, WSAudioSegment +from .ws_decode import decode_sample, decode_arr from .ws_sample import WSSample if TYPE_CHECKING: @@ -89,6 +89,8 @@ def get_sample(self, column: str, offset: int) -> typing.Any: try: if pa.types.is_binary(col_type) or pa.types.is_large_binary(col_type): return decode_sample(column, data) + if (column.rsplit(".", 1)[-1] if "." in column else "") == "arr": + return decode_arr(data, col_type) # native variable-length array return data.as_py(maps_as_pydicts="strict") except Exception as e: raise ValueError(f"Failed to decode column {column} in shard {self.fname} (offset {offset}): {e}") @@ -125,7 +127,7 @@ class WSSourceAudioShard(WSShardInterface): # cache _source_file_name: str = None _source_sample: WSSample = None - _source_reader: AudioReader = None + _source_reader: WSAudioEpisode = None @classmethod def from_link(cls, link, dataset, shard_ref): @@ -149,7 +151,7 @@ def get_sample(self, _column, offset): self._source_file_name = file_name tstart, tend = self.get_timestamps(segment_offset) - return WSAudio(self._source_reader, tstart, tend) + return WSAudioSegment(self._source_reader, tstart, tend) class WSYoutubeVideoShard(WSSourceAudioShard): diff --git a/wsds/ws_sink.py b/wsds/ws_sink.py index fb83cbb..3fca571 100644 --- a/wsds/ws_sink.py +++ b/wsds/ws_sink.py @@ -87,12 +87,41 @@ def write(self, x): if len(self._buffer) >= self.batch_size: self.write_batch(self._buffer) + def _build_record(self, b): + """Build a RecordBatch from the buffered rows. ".arr" columns are written + as native pyarrow list>/list (variable-length + arrays); everything else keeps the existing from_pylist inference.""" + import pyarrow + import numpy as np + from .ws_decode import arr_column_type + + names = list(b[0].keys()) if b else [] + arr_cols = [k for k in names if isinstance(k, str) + and k.rsplit(".", 1)[-1] == "arr"] + if not arr_cols: + return pyarrow.RecordBatch.from_pylist(b, self._sink_schema) + + arrays = [] + for name in names: + vals = [row.get(name) for row in b] + if name in arr_cols: + first = next((v for v in vals if v is not None), None) + if first is None: + arrays.append(pyarrow.array(vals)) + continue + t = arr_column_type(first) + arrays.append(pyarrow.array( + [None if v is None else np.asarray(v).tolist() for v in vals], type=t)) + else: + arrays.append(pyarrow.array(vals)) + return pyarrow.RecordBatch.from_arrays(arrays, names=names) + # TODO: test writing batches of data straight from a PyTorch batched processing loop def write_batch(self, b, flush=False): import pyarrow try: - record = pyarrow.RecordBatch.from_pylist(b, self._sink_schema) + record = self._build_record(b) except Exception: def _truncate(v, limit=200): r = repr(v)