|
from typing import Callable, Optional, Union |
|
from tqdm import tqdm |
|
import os |
|
import torch |
|
import torchaudio |
|
import torchaudio.functional as F |
|
from torch.utils.data import Dataset, DataLoader, IterableDataset, random_split |
|
from pytorch_lightning import LightningDataModule |
|
import webdataset |
|
|
|
|
|
class VLSP2020Dataset(Dataset): |
|
def __init__(self, root: str, sample_rate: int = 16000): |
|
super().__init__() |
|
|
|
self.sample_rate = sample_rate |
|
self.memory = self._prepare_data(root) |
|
self._memory = tuple( |
|
(v["transcript"], v["audio"]) for v in self.memory.values() |
|
) |
|
|
|
@staticmethod |
|
def _prepare_data(root: str): |
|
memory = {} |
|
|
|
for f in os.scandir(root): |
|
file_name, file_ext = os.path.splitext(f.name) |
|
|
|
if file_ext == ".txt": |
|
if file_name not in memory: |
|
memory[file_name] = {"transcript": f.path} |
|
elif "transcript" not in memory[file_name]: |
|
memory[file_name]["transcript"] = f.path |
|
else: |
|
raise ValueError(f"Duplicate transcript for {f.path}") |
|
else: |
|
if file_name not in memory: |
|
memory[file_name] = {"audio": f.path} |
|
elif "audio" not in memory[file_name]: |
|
memory[file_name]["audio"] = f.path |
|
else: |
|
raise ValueError(f"Duplicate audio for {f.path}") |
|
|
|
for key, value in memory.items(): |
|
if "audio" not in value: |
|
raise ValueError(f"Missing audio for {key}") |
|
elif "transcript" not in value: |
|
raise ValueError(f"Missing transcript for {key}") |
|
|
|
return memory |
|
|
|
def __len__(self): |
|
return len(self.memory) |
|
|
|
def __getitem__(self, index: int): |
|
transcript, audio = self._memory[index] |
|
|
|
with open(transcript, "r") as f: |
|
transcript = f.read() |
|
|
|
audio, sample_rate = torchaudio.load(audio) |
|
audio = F.resample(audio, sample_rate, self.sample_rate) |
|
|
|
return transcript, audio |
|
|
|
|
|
class VLSP2020TarDataset: |
|
def __init__(self, outpath: str): |
|
self.outpath = outpath |
|
|
|
def convert(self, dataset: VLSP2020Dataset): |
|
writer = webdataset.TarWriter(self.outpath) |
|
|
|
for idx, (transcript, audio) in enumerate(tqdm(dataset, colour="green")): |
|
writer.write( |
|
{ |
|
"__key__": f"{idx:08d}", |
|
"txt": transcript, |
|
"pth": audio, |
|
} |
|
) |
|
|
|
writer.close() |
|
|
|
def load(self) -> webdataset.WebDataset: |
|
self.data = ( |
|
webdataset.WebDataset(self.outpath) |
|
.decode( |
|
webdataset.handle_extension("txt", lambda x: x.decode("utf-8")), |
|
webdataset.torch_audio, |
|
) |
|
.to_tuple("txt", "pth") |
|
) |
|
|
|
return self.data |
|
|
|
|
|
def get_dataloader( |
|
dataset: Union[VLSP2020Dataset, webdataset.WebDataset], |
|
return_transcript: bool = False, |
|
target_transform: Optional[Callable] = None, |
|
batch_size: int = 32, |
|
num_workers: int = 2, |
|
): |
|
def collate_fn(batch): |
|
def get_audio(item): |
|
audio = item[1] |
|
|
|
assert ( |
|
isinstance(audio, torch.Tensor) |
|
and audio.ndim == 2 |
|
and audio.size(0) == 1 |
|
) |
|
|
|
return audio.squeeze(0) |
|
|
|
audio = tuple(get_audio(item) for item in batch) |
|
|
|
if return_transcript: |
|
if target_transform is not None: |
|
transcript = tuple(target_transform(item[0]) for item in batch) |
|
else: |
|
transcript = tuple(item[0] for item in batch) |
|
|
|
return transcript, audio |
|
else: |
|
return audio |
|
|
|
return DataLoader( |
|
dataset, batch_size=batch_size, num_workers=num_workers, collate_fn=collate_fn |
|
) |
|
|
