158 lines
7.5 KiB
Python
158 lines
7.5 KiB
Python
import os
|
|
import torch
|
|
import random
|
|
import copy
|
|
import logging
|
|
import shutil
|
|
|
|
class dataset(torch.utils.data.Dataset):
|
|
def __init__(self, args, split):
|
|
super().__init__()
|
|
self.args = args
|
|
self.split = split
|
|
assert self.split in ['train', 'validation', 'test']
|
|
manifest_fn = os.path.join(self.args.dataset_dir, self.args.manifest_name, self.split+".txt")
|
|
|
|
with open(manifest_fn, "r") as rf:
|
|
data = [l.strip().split("\t") for l in rf.readlines()]
|
|
lengths_list = [int(item[-1]) for item in data]
|
|
self.data = []
|
|
self.lengths_list = []
|
|
for d, l in zip(data, lengths_list):
|
|
if l >= self.args.encodec_sr*self.args.audio_min_length:
|
|
if self.args.drop_long and l > self.args.encodec_sr*self.args.audio_max_length:
|
|
continue
|
|
self.data.append(d)
|
|
self.lengths_list.append(l)
|
|
logging.info(f"number of data points for {self.split} split: {len(self.lengths_list)}")
|
|
|
|
# phoneme vocabulary
|
|
vocab_fn = os.path.join(self.args.dataset_dir,"vocab.txt")
|
|
shutil.copy(vocab_fn, os.path.join(self.args.exp_dir, "vocab.txt"))
|
|
with open(vocab_fn, "r") as f:
|
|
temp = [l.strip().split(" ") for l in f.readlines() if len(l) != 0]
|
|
self.phn2num = {item[1]:int(item[0]) for item in temp}
|
|
|
|
self.symbol_set = set(["<SIL>", "<MUSIC>", "<NOISE>", "<OTHER>"])
|
|
|
|
def __len__(self):
|
|
return len(self.lengths_list)
|
|
|
|
def _load_phn_enc(self, index):
|
|
item = self.data[index]
|
|
pf = os.path.join(self.args.dataset_dir, self.args.phn_folder_name, item[1]+".txt")
|
|
ef = os.path.join(self.args.dataset_dir, self.args.encodec_folder_name, item[1]+".txt")
|
|
try:
|
|
with open(pf, "r") as p, open(ef, "r") as e:
|
|
phns = [l.strip() for l in p.readlines()]
|
|
assert len(phns) == 1, phns
|
|
x = [self.phn2num[item] for item in phns[0].split(" ") if item not in self.symbol_set] # drop ["<SIL>", "<MUSIC>", "<NOISE>", "<OTHER>"], as they are not in training set annotation
|
|
encos = [l.strip().split() for k, l in enumerate(e.readlines()) if k < self.args.n_codebooks]
|
|
|
|
assert len(encos) == self.args.n_codebooks, ef
|
|
if self.args.special_first:
|
|
y = [[int(n)+self.args.n_special for n in l] for l in encos]
|
|
else:
|
|
y = [[int(n) for n in l] for l in encos]
|
|
if self.args.training_stage == 1 and not self.args.valle and not (self.args.musicgen or self.args.valle_orig):
|
|
y = y[:1]
|
|
except Exception as e:
|
|
logging.info(f"loading failed for {pf} and {ef}, maybe files don't exist or are corrupted")
|
|
logging.info(f"error message: {e}")
|
|
return [], [[]]
|
|
|
|
return x, y
|
|
|
|
def __getitem__(self, index):
|
|
x, y = self._load_phn_enc(index)
|
|
x_len, y_len = len(x), len(y[0])
|
|
|
|
if x_len == 0 or y_len == 0:
|
|
return {
|
|
"x": None,
|
|
"x_len": None,
|
|
"y": None,
|
|
"y_len": None,
|
|
"y_mask_interval": None, # index y_mask_interval[1] is the position of start_of_continue token
|
|
"extra_mask_start": None # this is only used in VE1
|
|
}
|
|
while y_len < self.args.encodec_sr*self.args.audio_min_length:
|
|
assert not self.args.dynamic_batching
|
|
index = random.choice(range(len(self))) # regenerate an index
|
|
x, y = self._load_phn_enc(index)
|
|
x_len, y_len = len(x), len(y[0])
|
|
if self.args.drop_long:
|
|
while x_len > self.args.text_max_length or y_len > self.args.encodec_sr*self.args.audio_max_length:
|
|
index = random.choice(range(len(self))) # regenerate an index
|
|
x, y = self._load_phn_enc(index)
|
|
x_len, y_len = len(x), len(y[0])
|
|
|
|
### padding and cropping below ###
|
|
### padding and cropping below ###
|
|
# adjust the length of encodec codes, pad to max_len or randomly crop
|
|
orig_y_len = copy.copy(y_len)
|
|
max_len = int(self.args.audio_max_length * self.args.encodec_sr)
|
|
if y_len > max_len:
|
|
audio_start = random.choice(range(0, y_len-max_len))
|
|
for i in range(len(y)):
|
|
y[i] = y[i][audio_start:(audio_start+max_len)]
|
|
y_len = max_len
|
|
else:
|
|
audio_start = 0
|
|
if not self.args.dynamic_batching:
|
|
pad = [0] * (max_len - y_len) if self.args.sep_special_token else [self.args.audio_pad_token] * (max_len - y_len)
|
|
for i in range(len(y)):
|
|
y[i] = y[i] + pad
|
|
|
|
# adjust text
|
|
# if audio is cropped, and text is longer than max, crop max based on how audio is cropped
|
|
if audio_start > 0 and len(x) > self.args.text_max_length: # if audio is longer than max and text is long than max, start text the way audio started
|
|
x = x[int(len(x)*audio_start/orig_y_len):]
|
|
if len(x) > self.args.text_max_length: # if text is still longer than max, cut the end
|
|
x = x[:self.args.text_max_length]
|
|
|
|
x_len = len(x)
|
|
if x_len > self.args.text_max_length:
|
|
text_start = random.choice(range(0, x_len - self.args.text_max_length))
|
|
x = x[text_start:text_start+self.args.text_max_length]
|
|
x_len = self.args.text_max_length
|
|
elif self.args.pad_x and x_len <= self.args.text_max_length:
|
|
pad = [0] * (self.args.text_max_length - x_len) if self.args.sep_special_token else [self.args.text_pad_token] * (self.args.text_max_length - x_len)
|
|
x = x + pad
|
|
### padding and cropping above ###
|
|
### padding and cropping above ###
|
|
|
|
return {
|
|
"x": torch.LongTensor(x),
|
|
"x_len": x_len,
|
|
"y": torch.LongTensor(y),
|
|
"y_len": y_len
|
|
}
|
|
|
|
|
|
def collate(self, batch):
|
|
out = {key:[] for key in batch[0]}
|
|
for item in batch:
|
|
if item['x'] == None: # deal with load failure
|
|
continue
|
|
for key, val in item.items():
|
|
out[key].append(val)
|
|
res = {}
|
|
if self.args.pad_x:
|
|
res["x"] = torch.stack(out["x"], dim=0)
|
|
else:
|
|
res["x"] = torch.nn.utils.rnn.pad_sequence(out["x"], batch_first=True, padding_value=0 if self.args.sep_special_token else self.args.text_pad_token)
|
|
res["x_lens"] = torch.LongTensor(out["x_len"])
|
|
if self.args.dynamic_batching:
|
|
if out['y'][0].ndim==2:
|
|
res['y'] = torch.nn.utils.rnn.pad_sequence([item.transpose(1,0) for item in out['y']],padding_value=0 if self.args.sep_special_token else self.args.audio_pad_token)
|
|
res['y'] = res['y'].permute(1,2,0) # T B K -> B K T
|
|
else:
|
|
assert out['y'][0].ndim==1, out['y'][0].shape
|
|
res['y'] = torch.nn.utils.rnn.pad_sequence(out['y'], batch_first=True, padding_value=0 if self.args.sep_special_token else self.args.audio_pad_token)
|
|
else:
|
|
res['y'] = torch.stack(out['y'], dim=0)
|
|
res["y_lens"] = torch.LongTensor(out["y_len"])
|
|
res["text_padding_mask"] = torch.arange(res['x'][0].shape[-1]).unsqueeze(0) >= res['x_lens'].unsqueeze(1)
|
|
res["audio_padding_mask"] = torch.arange(res['y'][0].shape[-1]).unsqueeze(0) >= res['y_lens'].unsqueeze(1)
|
|
return res |