373 lines
13 KiB
Python
373 lines
13 KiB
Python
import re
|
||
import sys
|
||
import json
|
||
from typing import List, Union
|
||
|
||
from torch.utils.data import Dataset
|
||
import torchaudio
|
||
from torchaudio.functional import resample
|
||
import torch
|
||
import numpy as np
|
||
|
||
from torch.nn.utils.rnn import pad_sequence
|
||
|
||
PARAGRAPH_GAP = 6
|
||
MIN_MUSIC_LEN = 3
|
||
|
||
def check_lryics(lyric):
|
||
_FILTER_STRING = [
|
||
'作词', '作曲', '编曲', '【', '策划',
|
||
'录音', '混音', '母带', '：', '制作',
|
||
'版权', '校对', '演奏', '制作', '伴奏'
|
||
]
|
||
for item in _FILTER_STRING:
|
||
if item in lyric:
|
||
return True
|
||
|
||
return False
|
||
|
||
|
||
|
||
def process_lyrics(lines):
|
||
lyric_part = []
|
||
timestamp_part = []
|
||
|
||
timestamp_pattern = re.compile(r'\[\d+:\d+(\.\d+)?\]')
|
||
|
||
for i, line in enumerate(lines):
|
||
|
||
# 删除前几行的特定信息
|
||
if i<10 and check_lryics(line):
|
||
continue
|
||
|
||
# 检查是否包含有效的时间戳和歌词内容
|
||
if timestamp_pattern.match(line):
|
||
timestamp_end = line.rfind(']')
|
||
lyrics = line[timestamp_end + 1:].strip()
|
||
timestamps = line[:timestamp_end + 1]
|
||
|
||
if '：' in lyrics:
|
||
if len(lyrics.split("：")[0]) <=5:
|
||
lyrics = "".join(lyrics.split("：")[1:])
|
||
# if lyrics: # 确保歌词部分不是空的
|
||
# lyric_part.append(lyrics)
|
||
# timestamp_part.append(timestamps)
|
||
# print(processed_lyrics)
|
||
return timestamp_part, lyric_part
|
||
|
||
def get_timestamps(timestamp_part):
|
||
|
||
# 转换为秒
|
||
|
||
timestamps = []
|
||
|
||
for line in timestamp_part:
|
||
match = re.match(r'\[(\d+):(\d+)(\.\d+)?\]', line)
|
||
if match:
|
||
minutes = int(match.group(1))
|
||
seconds = float(match.group(2))
|
||
millis = float(match.group(3)) if match.group(3) else 0
|
||
total_seconds = minutes * 60 + seconds + millis
|
||
timestamps.append(total_seconds)
|
||
|
||
|
||
return timestamps
|
||
|
||
def process_lyrics_lrc(lyrics):
|
||
timestamp_part, lyric_part = process_lyrics(lyrics)
|
||
# print(timestamp_part)
|
||
# print(lyric_part)
|
||
timestamps = get_timestamps(timestamp_part)
|
||
# print(timestamps)
|
||
if len(timestamps) == 0:
|
||
# print(f'{lyric_path}')
|
||
return []
|
||
|
||
slice_start = timestamps[0]
|
||
slice_start_idx = 0
|
||
|
||
output_list = []
|
||
for i in range(1, len(timestamps)):
|
||
# 如果累积时间超过30秒，则进行切分, 如果整体小于30s, 整句会被丢掉
|
||
if timestamps[i] - slice_start > 30:
|
||
output_list.append(f'[{str(slice_start)}:{str(timestamps[i])}]' + ", ".join(lyric_part[slice_start_idx:i]))
|
||
|
||
slice_start = timestamps[i]
|
||
slice_start_idx = i
|
||
|
||
return output_list
|
||
|
||
|
||
|
||
def process_lyrics_yrc(lyrics):
|
||
|
||
timestamps, lyric_part = extract_lrc(lyrics)
|
||
|
||
# timestamp_part, lyric_part = process_lyrics(lyrics)
|
||
# import pdb; pdb.set_trace()
|
||
# print(timestamp_part)
|
||
# print(lyric_part)
|
||
# timestamps = get_timestamps(timestamp_part)
|
||
# print(timestamps)
|
||
if len(timestamps) == 0:
|
||
# print(f'{lyric_path}')
|
||
return []
|
||
|
||
slice_start = timestamps[0]
|
||
slice_start_idx = 0
|
||
|
||
output_list = []
|
||
for i in range(1, len(timestamps)):
|
||
# 如果累积时间超过30秒，则进行切分
|
||
if timestamps[i] - slice_start > 30:
|
||
output_list.append(f'[{str(slice_start)}:{str(timestamps[i])}]' + ", ".join(lyric_part[slice_start_idx:i]))
|
||
|
||
slice_start = timestamps[i]
|
||
slice_start_idx = i
|
||
# import pdb; pdb.set_trace()
|
||
return output_list
|
||
|
||
def extract_lrc(lyrics):
|
||
timestamp_part, lyric_part = [], []
|
||
|
||
for i, text in enumerate(lyrics):
|
||
# 提取中括号内的内容
|
||
bracket_content = re.search(r'\[(.*?)\]', text).group(1)
|
||
bracket_content = bracket_content.split(',')
|
||
# 提取小括号内的内容
|
||
parentheses_content = re.findall(r'\((.*?)\)', text)
|
||
# 提取其他内容
|
||
other_content = re.sub(r'\[(.*?)\]|\((.*?)\)', '', text).strip()
|
||
|
||
# 数据怎么处理？
|
||
if i<10 and check_lryics(other_content):
|
||
continue
|
||
timestamp_part.append(float(bracket_content[0])/1000)
|
||
lyric_part.append(other_content)
|
||
return timestamp_part, lyric_part
|
||
|
||
|
||
|
||
class WYYSongDataset(Dataset):
|
||
def __init__(self,
|
||
metadata_path: Union[str, List[str]],
|
||
sr:int = 0,
|
||
use_lang = ['en', 'zh-cn'],
|
||
num_examples = -1,
|
||
max_dur = 20,
|
||
min_dur=0,
|
||
add_music=False,
|
||
pad_to_max= True,
|
||
):
|
||
|
||
self.sr = sr
|
||
self.use_lang = use_lang
|
||
self.data = []
|
||
if type(metadata_path) == str:
|
||
metadata_path = [metadata_path]
|
||
for _meta in metadata_path:
|
||
self._load_metadata(_meta)
|
||
self.max_dur = max_dur
|
||
self.min_dur = min_dur
|
||
self.pad_to_max = pad_to_max
|
||
self.add_music = add_music
|
||
|
||
# buffer
|
||
self.lyric_buffer = {}
|
||
|
||
if(num_examples<=0):
|
||
self.dataset_len = len(self.data)
|
||
self.random_slc = False
|
||
else:
|
||
self.dataset_len = num_examples
|
||
self.random_slc = True
|
||
|
||
|
||
# 读取jsonl文件
|
||
def _load_metadata(self, metadata_path):
|
||
with open(metadata_path) as fp:
|
||
lines = fp.readlines()
|
||
for line in lines:
|
||
item = json.loads(line)
|
||
if '伴奏' not in item['path']:
|
||
# if "lang_type" in item and item['lang_type'] == 'en':
|
||
if "lang_type" in item:
|
||
self.data.append(item)
|
||
|
||
|
||
def __len__(self):
|
||
return self.dataset_len
|
||
|
||
|
||
def __getitem__(self, idx):
|
||
try_cnt = 0
|
||
while True:
|
||
if(self.random_slc):
|
||
idx = np.random.randint(0, len(self.data))
|
||
yrc_lyrics = []
|
||
lrc_lyrics = []
|
||
try:
|
||
info = self.data[idx]
|
||
|
||
# audio path
|
||
path = info["path"]
|
||
lang_type = info["lang_type"]
|
||
lyrics = info['lyrics'] # chinese
|
||
# lyrics = info['lyrics_phone']
|
||
|
||
# 随机选取一个lyric段落
|
||
|
||
parsed_lyrics = []
|
||
# st_idx = np.random.randint(0, len(lyrics))
|
||
for ly_id in range(len(lyrics)):
|
||
lyric = lyrics[ly_id].strip()
|
||
st, et, lyric = self.parse_lyric(lyric)
|
||
|
||
if et - st >= self.max_dur:
|
||
continue #TODO 前后外沿 [MUSIC]
|
||
|
||
if parsed_lyrics != []:
|
||
if st - parsed_lyrics[-1][1] >= PARAGRAPH_GAP: # 大gap
|
||
parsed_lyrics.append((parsed_lyrics[-1][1], st, '[GAP]'))
|
||
elif self.add_music and st - parsed_lyrics[-1][1] >= MIN_MUSIC_LEN:
|
||
parsed_lyrics.append((parsed_lyrics[-1][1], st, '[MUSIC]'))
|
||
|
||
lyric = lyric.replace("\xa0", " ")
|
||
lyric = " ".join(lyric.split())
|
||
parsed_lyrics.append((st, et, lyric))
|
||
|
||
assert parsed_lyrics != []
|
||
# if parsed_lyrics[-1][1] - parsed_lyrics[0][0] > self.max_dur:
|
||
# print(f"{parsed_lyrics[0][0]}-{parsed_lyrics[-1][1]} {parsed_lyrics}", file=open('tmp.txt', 'a'))
|
||
|
||
parsed_lyrics = [(0, parsed_lyrics[0][0], '[GAP]')] + parsed_lyrics
|
||
|
||
possible_starts = [e for e,i in enumerate(parsed_lyrics) if i[2]=='[GAP]']
|
||
st_idx = np.random.choice(possible_starts)
|
||
|
||
paraphrase = []
|
||
for i in parsed_lyrics[st_idx+1:]:
|
||
if i[2] == '[GAP]':
|
||
break
|
||
paraphrase.append(i)
|
||
# print(paraphrase, lyrics)
|
||
|
||
while paraphrase[-1][1] - paraphrase[0][0] > self.max_dur:
|
||
if np.random.rand() > 0.2:
|
||
paraphrase.pop(-1) # 大概率从后面截断
|
||
else:
|
||
paraphrase.pop(0) # 小概率截前面
|
||
|
||
st, et, lyric = paraphrase[0][0], paraphrase[-1][1], ', '.join([i[2] for i in paraphrase]) # [SEP]
|
||
# print(st, et, lyric)
|
||
# import pdb; pdb.set_trace()
|
||
assert self.min_dur < et - st < self.max_dur, f"{st}-{et} {lyric}"
|
||
# print(et-st, lyric)
|
||
# import pdb; pdb.set_trace()
|
||
|
||
if info["lang_type"] == 'en':
|
||
# print(len(lyric.split())/(et-st))
|
||
char_num = sum([len(lrc[-1].split()) for lrc in paraphrase])
|
||
assert 6 > char_num / (et-st) > 1
|
||
else:
|
||
# print(len(lyric.split())/(et-st))
|
||
char_num = sum([len(lrc[-1]) for lrc in paraphrase])
|
||
assert 6 > char_num / (et-st) > 1
|
||
|
||
# 读取音频文件
|
||
cur_sample_rate = torchaudio.info(path).sample_rate
|
||
offset = int(cur_sample_rate*st)
|
||
num_frames = int(cur_sample_rate * (et -st))
|
||
chunk, _ = torchaudio.load(path, frame_offset=offset, num_frames=num_frames)
|
||
# chunk = torch.zeros(1, 48000*15)
|
||
if abs(chunk.shape[-1] - num_frames) > num_frames * 0.05: # 音频文件长度与歌词不一致
|
||
print(f"fail to load {path} from {st} to {et} !")
|
||
raise FileNotFoundError
|
||
# 随机选取一个channel
|
||
if(chunk.shape[0]>1):
|
||
chunk = chunk[torch.randint(chunk.shape[0], size=(1,)),:].float()
|
||
else:
|
||
chunk = chunk[[0],:].float()
|
||
|
||
if(cur_sample_rate!=self.sr):
|
||
# print('a:',cur_sample_rate,chunk.shape)
|
||
chunk = torchaudio.functional.resample(chunk, cur_sample_rate, self.sr)
|
||
|
||
if self.pad_to_max:
|
||
chunk = self.pad_2d_tensor(chunk, int(self.max_dur * self.sr), 0)
|
||
|
||
# print(self.sz_cnt)
|
||
return chunk, lyric, [st, et], path, lang_type
|
||
except (AssertionError, FileNotFoundError, RuntimeError) as e: # 其他Error不ok
|
||
# print("Error loadding ", info["path"])
|
||
try_cnt += 1
|
||
idx = np.random.randint(0, len(self.data))
|
||
if(try_cnt>100):
|
||
raise e
|
||
|
||
def parse_lyric(self, lyric):
|
||
pattern = r'\[(\d+\.\d+):(\d+\.\d+)\](.*)'
|
||
match = re.search(pattern, lyric)
|
||
|
||
start_time = float(match.group(1))
|
||
end_time = float(match.group(2))
|
||
content = match.group(3)
|
||
return start_time, end_time, content
|
||
|
||
def pad_2d_tensor(self, x, max_len, pad_id):
|
||
# 获取输入 tensor 的形状
|
||
batch_size, seq_len = x.size()
|
||
max_len = max(max_len, seq_len)
|
||
# 计算需要填充的长度
|
||
pad_len = max_len - seq_len
|
||
|
||
# 如果需要填充
|
||
if pad_len > 0:
|
||
# 创建填充 tensor
|
||
pad_tensor = torch.full((batch_size, pad_len), pad_id, dtype=x.dtype, device=x.device)
|
||
|
||
# 沿第二个维度（列）连接输入 tensor 和填充 tensor
|
||
padded_tensor = torch.cat([x, pad_tensor], dim=1)
|
||
else:
|
||
# 如果不需要填充，直接返回输入 tensor
|
||
padded_tensor = x
|
||
|
||
return padded_tensor
|
||
|
||
def collect_data(data_list):
|
||
audios = pad_sequence([data[0].t() for data in data_list], batch_first=True, padding_value=0).transpose(1,2)
|
||
lyrics = [data[1] for data in data_list]
|
||
st_et = [data[2] for data in data_list]
|
||
paths = [data[3] for data in data_list]
|
||
lang_types = [data[4] for data in data_list]
|
||
return audios, lyrics, st_et
|
||
# return audios, lyrics, st_et
|
||
|
||
|
||
def build_dataset(train_jsonl_list, val_jsonl_list, min_dur=0, max_dur=20, add_music=False):
|
||
print(min_dur,max_dur)
|
||
print(train_jsonl_list)
|
||
# ["exp/wyy3_20240418_v2f.jsonl",
|
||
# "exp/tme_lyric_baokuan.jsonl"]
|
||
train_dataset = WYYSongDataset(
|
||
metadata_path = train_jsonl_list,
|
||
sr = 48000,
|
||
use_lang = ['zh-cn', 'en'],
|
||
num_examples = 10*10000,
|
||
min_dur=min_dur,
|
||
max_dur=max_dur,
|
||
add_music=add_music
|
||
)
|
||
|
||
valid_dataset = WYYSongDataset(
|
||
metadata_path = val_jsonl_list,
|
||
sr = 48000,
|
||
use_lang = ['zh-cn', 'en'],
|
||
num_examples = 500,
|
||
min_dur=min_dur,
|
||
max_dur=max_dur,
|
||
add_music=add_music
|
||
)
|
||
print(train_jsonl_list, "\t total_song = ", len(train_dataset.data))
|
||
return train_dataset, valid_dataset
|