Refactor LibriSpeech Conformer RNN-T recipe

examples/asr/librispeech_conformer_rnnt/data_module.py

@@ -0,0 +1,191 @@
+import os
+import random
+
+import torch
+import torchaudio
+from pytorch_lightning import LightningDataModule
+
+
+def _batch_by_token_count(idx_target_lengths, max_tokens, batch_size=None):
+    batches = []
+    current_batch = []
+    current_token_count = 0
+    for idx, target_length in idx_target_lengths:
+        if current_token_count + target_length > max_tokens or (batch_size and len(current_batch) == batch_size):
+            batches.append(current_batch)
+            current_batch = [idx]
+            current_token_count = target_length
+        else:
+            current_batch.append(idx)
+            current_token_count += target_length
+
+    if current_batch:
+        batches.append(current_batch)
+
+    return batches
+
+
+def get_sample_lengths(librispeech_dataset):
+    fileid_to_target_length = {}
+
+    def _target_length(fileid):
+        if fileid not in fileid_to_target_length:
+            speaker_id, chapter_id, _ = fileid.split("-")
+
+            file_text = speaker_id + "-" + chapter_id + librispeech_dataset._ext_txt
+            file_text = os.path.join(librispeech_dataset._path, speaker_id, chapter_id, file_text)
+
+            with open(file_text) as ft:
+                for line in ft:
+                    fileid_text, transcript = line.strip().split(" ", 1)
+                    fileid_to_target_length[fileid_text] = len(transcript)
+
+        return fileid_to_target_length[fileid]
+
+    return [_target_length(fileid) for fileid in librispeech_dataset._walker]
+
+
+class CustomBucketDataset(torch.utils.data.Dataset):
+    def __init__(
+        self,
+        dataset,
+        lengths,
+        max_tokens,
+        num_buckets,
+        shuffle=False,
+        batch_size=None,
+    ):
+        super().__init__()
+
+        assert len(dataset) == len(lengths)
+
+        self.dataset = dataset
+
+        max_length = max(lengths)
+        min_length = min(lengths)
+
+        assert max_tokens >= max_length
+
+        buckets = torch.linspace(min_length, max_length, num_buckets)
+        lengths = torch.tensor(lengths)
+        bucket_assignments = torch.bucketize(lengths, buckets)
+
+        idx_length_buckets = [(idx, length, bucket_assignments[idx]) for idx, length in enumerate(lengths)]
+        if shuffle:
+            idx_length_buckets = random.sample(idx_length_buckets, len(idx_length_buckets))
+        else:
+            idx_length_buckets = sorted(idx_length_buckets, key=lambda x: x[1], reverse=True)
+
+        sorted_idx_length_buckets = sorted(idx_length_buckets, key=lambda x: x[2])
+        self.batches = _batch_by_token_count(
+            [(idx, length) for idx, length, _ in sorted_idx_length_buckets],
+            max_tokens,
+            batch_size=batch_size,
+        )
+
+    def __getitem__(self, idx):
+        return [self.dataset[subidx] for subidx in self.batches[idx]]
+
+    def __len__(self):
+        return len(self.batches)
+
+
+class TransformDataset(torch.utils.data.Dataset):
+    def __init__(self, dataset, transform_fn):
+        self.dataset = dataset
+        self.transform_fn = transform_fn
+
+    def __getitem__(self, idx):
+        return self.transform_fn(self.dataset[idx])
+
+    def __len__(self):
+        return len(self.dataset)
+
+
+class LibriSpeechDataModule(LightningDataModule):
+    def __init__(
+        self,
+        *,
+        librispeech_path,
+        train_transform,
+        val_transform,
+        test_transform,
+        max_tokens=700,
+        batch_size=2,
+        train_num_buckets=50,
+        train_shuffle=True,
+        num_workers=10,
+    ):
+        self.librispeech_path = librispeech_path
+        self.train_dataset_lengths = None
+        self.val_dataset_lengths = None
+        self.train_transform = train_transform
+        self.val_transform = val_transform
+        self.test_transform = test_transform
+        self.max_tokens = max_tokens
+        self.batch_size = batch_size
+        self.train_num_buckets = train_num_buckets
+        self.train_shuffle = train_shuffle
+        self.num_workers = num_workers
+
+    def train_dataloader(self):
+        datasets = [
+            torchaudio.datasets.LIBRISPEECH(self.librispeech_path, url="train-clean-360"),
+            torchaudio.datasets.LIBRISPEECH(self.librispeech_path, url="train-clean-100"),
+            torchaudio.datasets.LIBRISPEECH(self.librispeech_path, url="train-other-500"),
+        ]
+
+        if not self.train_dataset_lengths:
+            self.train_dataset_lengths = [get_sample_lengths(dataset) for dataset in datasets]
+
+        dataset = torch.utils.data.ConcatDataset(
+            [
+                CustomBucketDataset(
+                    dataset,
+                    lengths,
+                    self.max_tokens,
+                    self.train_num_buckets,
+                    batch_size=self.batch_size,
+                )
+                for dataset, lengths in zip(datasets, self.train_dataset_lengths)
+            ]
+        )
+        dataset = TransformDataset(dataset, self.train_transform)
+        dataloader = torch.utils.data.DataLoader(
+            dataset,
+            num_workers=self.num_workers,
+            batch_size=None,
+            shuffle=self.train_shuffle,
+        )
+        return dataloader
+
+    def val_dataloader(self):
+        datasets = [
+            torchaudio.datasets.LIBRISPEECH(self.librispeech_path, url="dev-clean"),
+            torchaudio.datasets.LIBRISPEECH(self.librispeech_path, url="dev-other"),
+        ]
+
+        if not self.val_dataset_lengths:
+            self.val_dataset_lengths = [get_sample_lengths(dataset) for dataset in datasets]
+
+        dataset = torch.utils.data.ConcatDataset(
+            [
+                CustomBucketDataset(
+                    dataset,
+                    lengths,
+                    self.max_tokens,
+                    1,
+                    batch_size=self.batch_size,
+                )
+                for dataset, lengths in zip(datasets, self.val_dataset_lengths)
+            ]
+        )
+        dataset = TransformDataset(dataset, self.val_transform)
+        dataloader = torch.utils.data.DataLoader(dataset, batch_size=None, num_workers=self.num_workers)
+        return dataloader
+
+    def test_dataloader(self):
+        dataset = torchaudio.datasets.LIBRISPEECH(self.librispeech_path, url="test-clean")
+        dataset = TransformDataset(dataset, self.test_transform)
+        dataloader = torch.utils.data.DataLoader(dataset, batch_size=None)
+        return dataloader

examples/asr/librispeech_conformer_rnnt/eval.py

@@ -4,7 +4,7 @@
 
 import torch
 import torchaudio
-from lightning import ConformerRNNTModule
+from lightning import ConformerRNNTModule, get_data_module
 
 
 logger = logging.getLogger()
@@ -15,19 +15,15 @@ def compute_word_level_distance(seq1, seq2):
 
 
 def run_eval(args):
-    model = ConformerRNNTModule.load_from_checkpoint(
-        args.checkpoint_path,
-        librispeech_path=str(args.librispeech_path),
-        sp_model_path=str(args.sp_model_path),
-        global_stats_path=str(args.global_stats_path),
-    ).eval()
+    model = ConformerRNNTModule.load_from_checkpoint(args.checkpoint_path, sp_model_path=str(args.sp_model_path)).eval()
+    data_module = get_data_module(str(args.librispeech_path), str(args.global_stats_path), str(args.sp_model_path))
 
     if args.use_cuda:
         model = model.to(device="cuda")
 
     total_edit_distance = 0
     total_length = 0
-    dataloader = model.test_dataloader()
+    dataloader = data_module.test_dataloader()
     with torch.no_grad():
         for idx, (batch, sample) in enumerate(dataloader):
             actual = sample[0][2]
@@ -45,6 +41,7 @@ def cli_main():
         "--checkpoint-path",
         type=pathlib.Path,
         help="Path to checkpoint to use for evaluation.",
+        required=True,
     )
     parser.add_argument(
         "--global-stats-path",
@@ -56,11 +53,13 @@ def cli_main():
         "--librispeech-path",
         type=pathlib.Path,
         help="Path to LibriSpeech datasets.",
+        required=True,
     )
     parser.add_argument(
         "--sp-model-path",
         type=pathlib.Path,
         help="Path to SentencePiece model.",
+        required=True,
     )
     parser.add_argument(
         "--use-cuda",