Repository: lhotse-speech/lhotse

Issue:
Large audio files processing
Piotr, 
I want to do some transcription of large audio files for some data-prep and I'm wondering the best way to do the processing.
These files may be too long to keep the audio in memory at one time, i.e. may be better to process them streaming or at least in chunks if there are efficient ways to process long compressed audios in chunks.
Basically this is for decoding but we'll probably want to write the decoding output to a jsonl file and later process it with a script, e.g. to stitch the transcripts of the pieces of audios back together.
I'm not sure what you think is the best way?  Does not necessarily have to involve Lhotse; in general, less framework is a good thing if it's not needed.  Efficiency is good though, I'd rather not delay the main decoding process by having it uncompress files and extract features.


Discussion:
Recent versions of `torchaudio` added the support for streaming audio loading using their own ffmpeg Python bindings. Check out their `StreamReader` API: https://pytorch.org/audio/stable/tutorials/streamreader_basic_tutorial.html 

If we want to make it maximally efficient, there's a fairly elegant solution: I could add a special `IterableDataset` that'd be created from a RecordingSet or CutSet, and when iterated, it yields mini-batches (of constant batch size and shape) of audio chunks and their cut/recording IDs. The I/O would happen in dataloader workers, and the only requirement is that the decoding script can handle adding and ending new "sessions" and managing the decoding state between mini-batches (perhaps Sherpa server + client is a good solution but it becomes "more framework").
I don't understand, if the I/O happens in dataloader workers, why the decoding script has to deal with sessions.
I was just planning to decode the audio in overlapping chunks, and there is no interaction, so we could write it out with id's
and join it later on.  (Maybe not nice from user-interface perspective though).

Got it, I think you mentioned in another thread you want to use offline ASR for each chunk. In that case it is quite simple indeed with no state and sessions. The dataloader + iterable dataset idea looks like a good fit (especially with torchaudio stream reader since it would open each file just once and return N chunks from it in each mini batch). 
OK, cool.  Can you please provide some concrete pointers on what exactly one would have to do?  E.g. how to configure the chunking?  Is the idea that we'd figure out the file lengths manually and then create chunks as a CutSet?


Resolving patch:
diff --git a/lhotse/dataset/unsupervised.py b/lhotse/dataset/unsupervised.py
--- a/lhotse/dataset/unsupervised.py
+++ b/lhotse/dataset/unsupervised.py
@@ -1,14 +1,17 @@
+import math
 import warnings
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 import torch
+from torch.utils.data import IterableDataset
 
-from lhotse import validate
-from lhotse.audio import suppress_audio_loading_errors
+from lhotse import RecordingSet, Seconds, compute_num_samples, validate
+from lhotse.audio import suppress_audio_loading_errors, torchaudio_supports_ffmpeg
 from lhotse.augmentation import AugmentFn
 from lhotse.cut import CutSet
 from lhotse.dataset.collation import collate_audio, collate_features, collate_matrices
 from lhotse.features import FeatureExtractor
+from lhotse.features.kaldi.layers import _get_strided_batch_streaming
 
 
 class UnsupervisedDataset(torch.utils.data.Dataset):
@@ -118,3 +121,169 @@ def generate_cut(cuts: CutSet):
     def _validate(self, cuts: CutSet) -> None:
         validate(cuts)
         assert all(cut.has_recording for cut in cuts)
+
+
+class RecordingChunkIterableDataset(IterableDataset):
+    """
+    This dataset iterates over chunks of a recording, for each recording provided.
+    It supports setting a chunk_shift < chunk_size to run model predictions on
+    overlapping audio chunks.
+
+    The format of yielded items is the following::
+        {
+            "recording_id": str
+            "begin_time": tensor with dtype=float32 shape=(1,)
+            "end_time": tensor with dtype=float32 shape=(1,)
+            "audio": tensor with dtype=float32 shape=(chunk_size_in_samples,)
+        }
+
+    Unlike most other datasets in Lhotse, this dataset does not yield batched items,
+    and should be used like the following::
+
+        >>> recordings = RecordingSet.from_file("my-recordings.jsonl.gz")
+        ... dataset = RecordingChunkIterableDataset(recordings, chunk_size=30.0, chunk_shift=25.0)
+        ... dloader = torch.utils.data.DataLoader(
+        ...     dataset,
+        ...     batch_size=32,
+        ...     collate_fn=audio_chunk_collate,
+        ...     worker_init_fn=audio_chunk_worker_init_fn,
+        ... )
+
+    """
+
+    def __init__(
+        self, recordings: RecordingSet, chunk_size: Seconds, chunk_shift: Seconds
+    ) -> None:
+        self.recordings = list(recordings)
+        self.chunk_size = chunk_size
+        self.chunk_shift = chunk_shift
+        self.start = 0
+        self.end = len(self.recordings)
+
+        self.validate()
+
+    def validate(self) -> None:
+        if not torchaudio_supports_ffmpeg():
+            raise RuntimeError(
+                "Using FFMPEG streamer backend for reading is supported only "
+                "with PyTorch 1.12+ and torchaudio 0.12+"
+            )
+
+        for r in self.recordings:
+            assert (
+                len(r.sources) == 1
+            ), f"We currently don't support multi-source audio in this dataset (got {len(r.sources)} sources in recording {r.id})."
+            assert (
+                r.sources[0].type == "file"
+            ), f"We currently only support 'file' AudioSource type in this dataset (got: {r.sources[0].type} in recording {r.id})."
+            assert (
+                r.num_channels == 1
+            ), f"We currently only support single-channel audio in this dataset (got {r.num_channels} channels in recording {r.id})."
+
+    def __iter__(self):
+        import torchaudio
+
+        for r in self.recordings[self.start : self.end]:
+            chunk_size = compute_num_samples(self.chunk_size, r.sampling_rate)
+            chunk_shift = compute_num_samples(self.chunk_shift, r.sampling_rate)
+
+            streamer = torchaudio.io.StreamReader(src=r.sources[0].source)
+            assert streamer.num_src_streams == 1, (
+                "Lhotse doesn't support files with more than one FFMPEG source stream yet "
+                "(not to be confused with multi-channel)."
+            )
+            streamer.add_basic_audio_stream(frames_per_chunk=chunk_size)
+
+            begin_time = 0
+            end_time = self.chunk_size
+            buffer = ShiftingBuffer(chunk_size=chunk_size, chunk_shift=chunk_shift)
+            for (incoming_audio,) in streamer.stream():
+                buffer.push(incoming_audio.squeeze())
+                for chunk in buffer.get_chunks():
+                    yield {
+                        "recording_id": r.id,
+                        "begin_time": torch.as_tensor(begin_time, dtype=torch.float32),
+                        "end_time": torch.as_tensor(end_time, dtype=torch.float32),
+                        "audio": chunk,
+                    }
+                    begin_time += self.chunk_shift
+                    end_time = begin_time + self.chunk_size
+            remainder = buffer.flush()
+            if remainder.shape[0] > 0:
+                yield {
+                    "recording_id": r.id,
+                    "begin_time": torch.as_tensor(begin_time, dtype=torch.float32),
+                    "end_time": torch.as_tensor(end_time, dtype=torch.float32),
+                    "audio": remainder,
+                }
+
+
+class ShiftingBuffer:
+    """
+    Utility for iterating over streaming audio chunks that supports chunk_shift < chunk_size.
+    It is useful when running model predictions on overlapping chunks of audio data.
+    """
+
+    def __init__(self, chunk_size: int, chunk_shift: int):
+        self.buf = torch.empty(1, 0)
+        self.chunk_size = chunk_size
+        self.chunk_shift = chunk_shift
+
+    def push(self, audio: torch.Tensor) -> None:
+        """Add new chunk of audio to the buffer. Expects shape (num_samples, )."""
+        self.buf = torch.cat([self.buf, audio.unsqueeze(0)], dim=1)
+
+    def get_chunks(self) -> torch.Tensor:
+        """
+        Retrieve chunks accumulated so far, adjusted for chunk_shift.
+        For chunk_shift < chunk_size, there will typically be more chunks
+        returned from this function than were pushed into the buffer because
+        of overlap.
+        The returned shape is (num_chunks, chunk_size).
+        """
+        out, self.buf = _get_strided_batch_streaming(
+            self.buf,
+            window_shift=self.chunk_shift,
+            window_length=self.chunk_size,
+            snip_edges=True,
+        )
+        return out.squeeze(0)
+
+    def flush(self) -> torch.Tensor:
+        """
+        Flush out the remainder chunk from the buffer.
+        Typically it will be shorter than chunk_size.
+        The returned shape is (remainder_size, ).
+        """
+        out = self.buf.squeeze(0)
+        self.buf = torch.empty(1, 0)
+        return out
+
+
+def audio_chunk_collate(batch: List[Dict]):
+    from torch.utils.data import default_collate
+
+    audios = [d.pop("audio") for d in batch]
+    out = default_collate(batch)
+
+    maxlen = max(a.shape[0] for a in audios)
+    audio = torch.zeros((len(audios), maxlen))
+    for i, a in enumerate(audios):
+        audio[i, : a.shape[0]] = torch.as_tensor(a)
+    out["audio"] = audio
+
+    return out
+
+
+def audio_chunk_worker_init_fn(worker_id):
+    worker_info = torch.utils.data.get_worker_info()
+    dataset = worker_info.dataset  # the dataset copy in this worker process
+    overall_start = dataset.start
+    overall_end = dataset.end
+    # configure the dataset to only process the split workload
+    per_worker = int(
+        math.ceil((overall_end - overall_start) / float(worker_info.num_workers))
+    )
+    worker_id = worker_info.id
+    dataset.start = overall_start + worker_id * per_worker
+    dataset.end = min(dataset.start + per_worker, overall_end)
diff --git a/lhotse/features/kaldi/layers.py b/lhotse/features/kaldi/layers.py
--- a/lhotse/features/kaldi/layers.py
+++ b/lhotse/features/kaldi/layers.py
@@ -853,7 +853,7 @@ def _get_strided_batch_streaming(
 
     if snip_edges:
         if num_samples < window_length:
-            return torch.empty((batch_size, 0, 0)), torch.empty(batch_size, 0)
+            return torch.empty((batch_size, 0, 0)), waveform
         num_frames = 1 + (num_samples - window_length) // window_shift
     else:
         window_remainder = window_length - window_shift


Test patch:
diff --git a/test/dataset/test_audio_chunk_dataset.py b/test/dataset/test_audio_chunk_dataset.py
new file mode 100644
--- /dev/null
+++ b/test/dataset/test_audio_chunk_dataset.py
@@ -0,0 +1,67 @@
+import pytest
+import torch
+from torch.utils.data import DataLoader
+
+from lhotse import RecordingSet, Seconds, compute_num_samples
+from lhotse.audio import torchaudio_supports_ffmpeg
+from lhotse.dataset.unsupervised import (
+    RecordingChunkIterableDataset,
+    audio_chunk_collate,
+    audio_chunk_worker_init_fn,
+)
+
+
+@pytest.mark.skipif(
+    not torchaudio_supports_ffmpeg(), reason="Requires torchaudio 0.12.0+ to run."
+)
+@pytest.mark.parametrize("CHUNK_SHIFT", [10.0, 8.0])
+def test_audio_chunk_dataset_usage(CHUNK_SHIFT: Seconds):
+    CHUNK_SIZE: Seconds = 10.0
+    SAMPLING_RATE = 16000
+    EXPECTED_CHUNK_NUM_SAMPLES = compute_num_samples(CHUNK_SIZE, SAMPLING_RATE)
+
+    # RecordingSet of 4 audio files
+    recordings = RecordingSet.from_file("test/fixtures/libri/audio.json").repeat(4)
+
+    dataset = RecordingChunkIterableDataset(
+        recordings, chunk_size=CHUNK_SIZE, chunk_shift=CHUNK_SHIFT
+    )
+
+    dloader = DataLoader(
+        dataset,
+        batch_size=2,
+        collate_fn=audio_chunk_collate,
+        num_workers=0,
+        worker_init_fn=audio_chunk_worker_init_fn,
+    )
+
+    tot_items = 0
+    for batch in dloader:
+        assert isinstance(batch, dict)
+
+        assert set(batch.keys()) == {"recording_id", "begin_time", "end_time", "audio"}
+
+        assert isinstance(batch["recording_id"], list)
+        assert len(batch["recording_id"]) == 2
+        assert isinstance(batch["recording_id"][0], str)
+        assert isinstance(batch["recording_id"][1], str)
+
+        assert torch.is_tensor(batch["begin_time"])
+        assert batch["begin_time"].shape == (2,)
+        assert batch["begin_time"].dtype == torch.float
+
+        assert torch.is_tensor(batch["end_time"])
+        assert batch["end_time"].shape == (2,)
+        assert batch["end_time"].dtype == torch.float
+
+        assert torch.is_tensor(batch["audio"])
+        assert batch["audio"].shape == (2, EXPECTED_CHUNK_NUM_SAMPLES)
+        assert batch["audio"].dtype == torch.float
+
+        tot_items += batch["audio"].shape[0]
+
+    # Check that no items were duplicated in the dataloader.
+    EXPECTED_TOT_ITEMS = (
+        len(recordings) * 2
+    )  # == 8 (because there are 2 chunks per recording)
+    assert tot_items == EXPECTED_TOT_ITEMS