# Copyright 2023 (authors: Feiteng Li) # # 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. """ modified from lhoste.dataset.speech_synthesis.py """ import torch import math import h5py from tokenizers import Tokenizer from typing import Union, List import numpy as np from tqdm import tqdm from utils.g2p import PhonemeBpeTokenizer from data.collation import get_text_token_collater text_collater = get_text_token_collater() _pad = '_' _punctuation = ',.!?-~…' _letters = 'NQabdefghijklmnopstuvwxyzɑæʃʑçɯɪɔɛɹðəɫɥɸʊɾʒθβŋɦ⁼ʰ`^#*=ˈˌ→↓↑ ' symbols = [_pad] + list(_punctuation) + list(_letters) language_dict = { 'en': 0, 'zh': 1, 'ja': 2, 'vi': 3, } def seq2phone(tokens: Union[List, np.ndarray]): """ Convert tokenized phoneme ID sequence back to phoneme string :param tokens: phoneme tokens :return: recovered phoneme sequence """ phones = "".join([symbols[i] for i in tokens]) return phones class DynamicBatchSampler(torch.utils.data.Sampler): def __init__(self, sampler, num_tokens_fn, num_buckets=100, min_size=0, max_size=1000, max_tokens=None, max_sentences=None, drop_last=False): """ :param sampler: :param num_tokens_fn: 根据idx返回样本的长度的函数 :param num_buckets: 利用桶原理将相似长度的样本放在一个batchsize中，桶的数量 :param min_size: 最小长度的样本， 小于这个值的样本会被过滤掉。 依据这个值来创建样桶 :param max_size: 最大长度的样本 :param max_sentences: batch_size, 但是这里可以通过max_sentences 和 max_tokens 共同控制最终的大小 """ super(DynamicBatchSampler, self).__init__(sampler) self.sampler = sampler self.num_tokens_fn = num_tokens_fn self.num_buckets = num_buckets self.min_size = min_size self.max_size = max_size assert max_size <= max_tokens, "max_size should be smaller than max tokens" assert max_tokens is not None or max_sentences is not None, \ "max_tokens and max_sentences should not be null at the same time, please specify one parameter at least" self.max_tokens = max_tokens if max_tokens is not None else float('Inf') self.max_sentences = max_sentences if max_sentences is not None else float('Inf') self.drop_last = drop_last def set_epoch(self, epoch): self.sampler.set_epoch(epoch) def is_batch_full(self, num_tokens, batch): if len(batch) == 0: return False if len(batch) == self.max_sentences: return True if num_tokens > self.max_tokens: return True return False def __iter__(self): buckets = [[] for _ in range(self.num_buckets)] sample_len = [0] * self.num_buckets for idx in self.sampler: idx_length = self.num_tokens_fn(idx) if not (self.min_size <= idx_length <= self.max_size): print("sentence at index {} of size {} exceeds max_tokens, the sentence is ignored".format(idx, idx_length)) continue index_buckets = math.floor((idx_length - self.min_size) / (self.max_size - self.min_size + 1) * self.num_buckets) sample_len[index_buckets] = max(sample_len[index_buckets], idx_length) num_tokens = (len(buckets[index_buckets]) + 1) * sample_len[index_buckets] if self.is_batch_full(num_tokens, buckets[index_buckets]): # yield this batch yield buckets[index_buckets] buckets[index_buckets] = [] sample_len[index_buckets] = 0 buckets[index_buckets].append(idx) # process left-over leftover_batch = [] leftover_sample_len = 0 leftover = [idx for bucket in buckets for idx in bucket] for idx in leftover: idx_length = self.num_tokens_fn(idx) leftover_sample_len = max(leftover_sample_len, idx_length) num_tokens = (len(leftover_batch) + 1) * leftover_sample_len if self.is_batch_full(num_tokens, leftover_batch): yield leftover_batch leftover_batch = [] leftover_sample_len = 0 leftover_batch.append(idx) if len(leftover_batch) > 0 and not self.drop_last: yield leftover_batch def __len__(self): # we do not know the exactly batch size, so do not call len(dataloader) pass class AudioDataset(torch.utils.data.Dataset): def __init__(self, h5_path, ann_path, tokenizer_path): self.h5_path = h5_path with open(ann_path, 'r', encoding='utf-8') as f: lines = f.readlines() ls = [l.split("|") for l in lines] ls_T = list(zip(*ls)) #del ls_T[-1] self.h5_paths, self.durations, self.langs, self.texts = \ list(ls_T[0]), list(ls_T[1]), list(ls_T[2]), list(ls_T[3]) self.durations = [float(dur) for dur in self.durations] self.tokenizer = PhonemeBpeTokenizer(tokenizer_path) self._archive = None def __len__(self): return len(self.h5_paths) def get_dur(self, idx): return self.durations[idx] @property def archive(self): if self._archive is None: # lazy loading here! self._archive = h5py.File(self.h5_path, "r") return self._archive def __getitem__(self, idx): archive = self.archive h5_path = self.h5_paths[idx] sub = archive[h5_path] audio_tokens = sub['audio'][()] #phone_tokens = sub['text'][()] dur = self.durations[idx] lang = self.langs[idx] text = self.texts[idx] # tokenization should be done within dataloader #phones = seq2phone(phone_tokens) #phones = phones.replace(" ", "_") phonemes, langs = self.tokenizer.tokenize(text=f"{text}".strip()) cptpho_tokens, enroll_x_lens = text_collater([phonemes]) cptpho_tokens = cptpho_tokens.squeeze(0) text_token_lens = enroll_x_lens[0] ''' if not len(phones): cptpho_tokens = self.tokenizer.encode(text).ids else: cptpho_tokens = self.tokenizer.encode(phones).ids ''' assert len(cptpho_tokens) return { 'utt_id': h5_path, 'text': text, 'audio': None, 'audio_lens': None, 'audio_features': audio_tokens, 'audio_features_lens': audio_tokens.shape[1], 'text_tokens': np.array(cptpho_tokens), 'text_tokens_lens': text_token_lens, 'language': language_dict[lang], } def collate(batch): utt_id_s = [b['utt_id'] for b in batch] text_s = [b['text'] for b in batch] audio_s = [b['audio'] for b in batch] audio_lens_s = [b['audio_lens'] for b in batch] audio_features_lens_s = [b['audio_features_lens'] for b in batch] # create an empty tensor with maximum audio feature length audio_features_s = torch.zeros([len(batch), max(audio_features_lens_s), 8], dtype=torch.int64) - 1 # audio pad with -1 text_tokens_lens_s = [b['text_tokens_lens'] for b in batch] # create an empty tensor with maximum text tokens length text_tokens_s = torch.zeros([len(batch), max(text_tokens_lens_s)], dtype=torch.int64) + 3 # [PAD] token id 3 language_s = [b['language'] for b in batch] for i, b in enumerate(batch): audio_features = b['audio_features'] audio_features_lens = b['audio_features_lens'] audio_features_s[i, :audio_features_lens, :] = torch.LongTensor(audio_features) text_tokens = b['text_tokens'] text_tokens_lens = b['text_tokens_lens'] text_tokens_s[i, :text_tokens_lens] = torch.LongTensor(text_tokens) batch = { 'utt_id': utt_id_s, 'text': text_s, 'audio': audio_s, 'audio_lens': audio_lens_s, 'audio_features': audio_features_s, 'audio_features_lens': torch.LongTensor(np.array(audio_features_lens_s)), 'text_tokens': text_tokens_s, 'text_tokens_lens': torch.LongTensor(np.array(text_tokens_lens_s)), 'languages': torch.LongTensor(np.array(language_s)), } return batch def create_dataloader(data_dir="/root/valle/egs/mix", n_gpus=1, rank=0, num_workers=0, num_buckets=10, max_duration=120): train_dataset = AudioDataset(h5_path=f"{data_dir}/audio_sum.hdf5", ann_path=f"{data_dir}/audio_ann_sum.txt", tokenizer_path=f"{data_dir}/bpe_175.json") ran_sampler = torch.utils.data.distributed.DistributedSampler( train_dataset, num_replicas=n_gpus, rank=rank, shuffle=True, ) dynamic_sampler = DynamicBatchSampler(ran_sampler, train_dataset.get_dur, num_buckets=num_buckets, max_size=20, max_tokens=max_duration) train_loader = torch.utils.data.DataLoader(train_dataset, num_workers=num_workers, collate_fn=collate, batch_sampler=dynamic_sampler) return train_loader