Initial commit: FunASR Speech Recognition Toolkit

Add complete FunASR codebase including models, runtime, and documentation.

funasr/models/eend/eend_ola_dataloader.py

@@ -0,0 +1,86 @@
+import logging
+
+import kaldiio
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+from torch.utils.data import Dataset
+
+
+def custom_collate(batch):
+    """Custom collate.
+    
+        Args:
+            batch: TODO.
+        """
+    keys, speech, speaker_labels, orders = zip(*batch)
+    speech = [torch.from_numpy(np.copy(sph)).to(torch.float32) for sph in speech]
+    speaker_labels = [torch.from_numpy(np.copy(spk)).to(torch.float32) for spk in speaker_labels]
+    orders = [torch.from_numpy(np.copy(o)).to(torch.int64) for o in orders]
+    batch = dict(speech=speech, speaker_labels=speaker_labels, orders=orders)
+
+    return keys, batch
+
+
+class EENDOLADataset(Dataset):
+    def __init__(
+        self,
+        data_file,
+    ):
+        """Initialize EENDOLADataset.
+        
+            Args:
+                data_file: TODO.
+            """
+        self.data_file = data_file
+        with open(data_file) as f:
+            lines = f.readlines()
+        self.samples = [line.strip().split() for line in lines]
+        logging.info("total samples: {}".format(len(self.samples)))
+
+    def __len__(self):
+        """Internal: len  ."""
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        """Internal: getitem  .
+        
+            Args:
+                idx: TODO.
+            """
+        key, speech_path, speaker_label_path = self.samples[idx]
+        speech = kaldiio.load_mat(speech_path)
+        speaker_label = kaldiio.load_mat(speaker_label_path).reshape(speech.shape[0], -1)
+
+        order = np.arange(speech.shape[0])
+        np.random.shuffle(order)
+
+        return key, speech, speaker_label, order
+
+
+class EENDOLADataLoader:
+    def __init__(self, data_file, batch_size, shuffle=True, num_workers=8):
+        """Initialize EENDOLADataLoader.
+        
+            Args:
+                data_file: TODO.
+                batch_size: Number of samples per batch.
+                shuffle: TODO.
+                num_workers: TODO.
+            """
+        dataset = EENDOLADataset(data_file)
+        self.data_loader = DataLoader(
+            dataset,
+            batch_size=batch_size,
+            collate_fn=custom_collate,
+            shuffle=shuffle,
+            num_workers=num_workers,
+        )
+
+    def build_iter(self, epoch):
+        """Build iter.
+        
+            Args:
+                epoch: TODO.
+            """
+        return self.data_loader