codeclm/tokenizer/Flow1dVAE/libs/datasets/dataset_429.py

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