icefall/egs/gigaspeech2/SSL/zipformer/dataset.py

# Copyright      2024  Xiaomi Corporation        (authors: Yifan Yang)
#
# See ../LICENSE for clarification regarding multiple authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys
from typing import Any, Dict, Optional

import numpy as np
import torch
import torch.nn.functional as F
from lhotse import validate
from lhotse.cut import CutSet
from lhotse.dataset.collation import collate_features
from lhotse.workarounds import Hdf5MemoryIssueFix
from torch.utils.data.dataloader import default_collate


class HubertDataset(torch.utils.data.Dataset):
    """
    In this implementation, there will always be a single channel.

    Returns:

    .. code-block::

        {
            'features': (B, T, F) float tensor
        }

    Dimension symbols legend:
    * ``B`` - batch size (number of Cuts)
    * ``T`` - number of frames of the longest Cut
    * ``F`` - number of features
    """

    def __init__(
        self,
        max_sample_size: Optional[int] = None,
        sample_rate: float = 100,
        label_rate: float = 50,
        random_crop: bool = True,
        pad_audio: bool = False,
        num_classes: list = [504],
    ) -> None:
        super().__init__()
        self.sample_rate = sample_rate
        self.label_rate = label_rate
        self.random_crop = random_crop
        self.pad_feature = pad_audio
        self.num_classes = num_classes
        self.max_sample_size = (
            max_sample_size if max_sample_size is not None else sys.maxsize
        )

        # This attribute is a workaround to constantly growing HDF5 memory
        # throughout the epoch. It regularly closes open file handles to
        # reset the internal HDF5 caches.
        self.hdf5_fix = Hdf5MemoryIssueFix(reset_interval=100)

    def __getitem__(self, cuts: CutSet) -> Dict[str, Any]:
        self._validate(cuts)
        self.hdf5_fix.update()

        # Sort the cuts by duration so that the first one determines the batch time dimensions.
        cuts = cuts.sort_by_duration(ascending=False)

        features = [torch.from_numpy(cut.load_features()) for cut in cuts]
        feature_lens = [cut.num_frames for cut in cuts]

        if self.pad_feature:
            feature_size = min(max(feature_lens), self.max_sample_size)
        else:
            feature_size = min(min(feature_lens), self.max_sample_size)

        features, padding_mask, feature_starts = self.collater_feature(
            features, feature_lens, feature_size
        )

        kmeans = [cut.custom["kmeans"] for cut in cuts]
        kmeans = [
            torch.tensor([int(item) for item in label.split()], dtype=torch.int64)
            for label in kmeans
        ]
        kmeans, kmeans_lens = self.collater_frm_label(kmeans, feature_size, feature_starts)

        return {
            "cuts": cuts,
            "features": features,
            "padding_mask": padding_mask,
            "kmeans": kmeans,
        }

    def _validate(self, cuts: CutSet) -> None:
        validate(cuts)
        assert all(cut.has_recording for cut in cuts)

    def crop_to_max_size(self, feature, target_size):
        size = len(feature)
        diff = size - target_size
        if diff <= 0:
            return feature, 0

        start, end = 0, target_size
        if self.random_crop:
            start = np.random.randint(0, diff + 1)
            end = size - diff + start
        return feature[start:end, :], start

    def collater_feature(self, features, feature_lens, feature_size):
        feature_dim = features[0].shape[-1]
        collated_features = features[0].new_zeros(len(features), feature_size, feature_dim)
        padding_mask = (
            torch.BoolTensor(collated_features.shape[:-1]).fill_(False)
            # if self.pad_feature else None
        )
        feature_starts = [0 for _ in features]
        for i, (feature, feature_len) in enumerate(zip(features, feature_lens)):
            diff = feature_len - feature_size
            if diff == 0:
                collated_features[i] = feature
            elif diff < 0:
                assert self.pad_feature
                collated_features[i] = torch.cat([feature, feature.new_full((-diff, feature_dim), 0.0)])
                padding_mask[i, diff:] = True
            else:
                collated_features[i], feature_starts[i] = self.crop_to_max_size(
                    feature, feature_size
                )
        return collated_features, padding_mask, feature_starts

    def collate_tokens(
        self,
        values,
        pad_idx,
        eos_idx=None,
        left_pad=False,
        move_eos_to_beginning=False,
        pad_to_length=None,
        pad_to_multiple=1,
        pad_to_bsz=None,
    ):
        """Convert a list of 1d tensors into a padded 2d tensor."""
        size = max(v.size(0) for v in values)
        size = size if pad_to_length is None else max(size, pad_to_length)
        if pad_to_multiple != 1 and size % pad_to_multiple != 0:
            size = int(((size - 0.1) // pad_to_multiple + 1) * pad_to_multiple)

        batch_size = len(values) if pad_to_bsz is None else max(len(values), pad_to_bsz)
        res = values[0].new(batch_size, size).fill_(pad_idx)

        def copy_tensor(src, dst):
            assert dst.numel() == src.numel()
            if move_eos_to_beginning:
                if eos_idx is None:
                    # if no eos_idx is specified, then use the last token in src
                    dst[0] = src[-1]
                else:
                    dst[0] = eos_idx
                dst[1:] = src[:-1]
            else:
                dst.copy_(src)

        for i, v in enumerate(values):
            copy_tensor(v, res[i][size - len(v) :] if left_pad else res[i][: len(v)])
        return res

    def collater_frm_label(self, targets, feature_size, feature_starts):
        label_rate = self.label_rate
        pad = self.num_classes[0] - 1
        assert label_rate > 0
        s2f = label_rate / self.sample_rate
        frm_starts = [int(round(s * s2f)) for s in feature_starts]
        frm_size = int(round(feature_size * s2f))
        if not self.pad_feature:
            rem_size = [len(t) - s for t, s in zip(targets, frm_starts)]
            frm_size = min(frm_size, *rem_size)
        targets = [t[s : s + frm_size] for t, s in zip(targets, frm_starts)]

        lengths = torch.LongTensor([len(t) for t in targets])
        targets = self.collate_tokens(targets, pad_idx=pad, left_pad=False)
        return targets, lengths


if __name__ == "__main__":
    from lhotse import load_manifest_lazy
    from lhotse.dataset import DynamicBucketingSampler
    from torch.utils.data import DataLoader

    dataset = HubertDataset(max_sample_size=1562)
    cuts = load_manifest_lazy("data/fbank/librispeech_cuts_train-clean-100.jsonl.gz")
    sampler = DynamicBucketingSampler(
        cuts,
        max_duration=300,
        shuffle=False,
    )
    dl = DataLoader(
        dataset,
        batch_size=None,
        sampler=sampler,
        num_workers=0,
    )

    for batch_idx, batch in enumerate(dl):
        print(batch["features"].shape)
        print(batch["padding_mask"].shape)
        print(batch["kmeans"].shape)