import os import random import re import numpy as np from tokenizers import AddedToken import torch import torch.nn.functional as F from transformers import PreTrainedTokenizerFast # to avoid: "The current process just got forked, after parallelism has already been used" os.environ["TOKENIZERS_PARALLELISM"] = "False" global tokenizer g_tokenizer = None dataset_to_token = { "youtube_music": "[YTM]", "youtube_music_lyrics": "[YTML]", "youtube_music_lyrics_foreign": "[YTMLF]", "genius_hq_lyrics": "[GEN]", "genius_hq_lyrics_foreign": "[GENF]", "imslp": "[IMSLP]", "jamendo": "[MJ]", "pond5_music": "[PD5]", "ytm_tagged": "[YTMT]", } BLACK_LIST = { "low rolloff", "high rolloff", "stage screen", "musicbeds", "classic hits", } def _load_tokenizer(tokenizer_fp=None): global g_tokenizer if g_tokenizer is not None: return g_tokenizer assert os.path.exists(tokenizer_fp) g_tokenizer = PreTrainedTokenizerFast( tokenizer_file=tokenizer_fp, unk_token="[UNK]", pad_token="[PAD]", ) g_tokenizer.add_special_tokens({"additional_special_tokens": [AddedToken("\n")]}) g_tokenizer.add_special_tokens({"additional_special_tokens": [AddedToken("[EOS]")]}) g_tokenizer.add_special_tokens({"additional_special_tokens": [AddedToken("[BOS]")]}) g_tokenizer.add_special_tokens({"additional_special_tokens": [AddedToken("[BOSL]")]}) g_tokenizer.add_special_tokens({"additional_special_tokens": [AddedToken("[YTM]")]}) g_tokenizer.add_special_tokens({"additional_special_tokens": [AddedToken("[YTML]")]}) g_tokenizer.add_special_tokens({"additional_special_tokens": [AddedToken("[YTMLF]")]}) g_tokenizer.add_special_tokens({"additional_special_tokens": [AddedToken("[GEN]")]}) g_tokenizer.add_special_tokens({"additional_special_tokens": [AddedToken("[GENF]")]}) g_tokenizer.add_special_tokens({"additional_special_tokens": [AddedToken("[IMSLP]")]}) g_tokenizer.add_special_tokens({"additional_special_tokens": [AddedToken("[MJ]")]}) g_tokenizer.add_special_tokens({"additional_special_tokens": [AddedToken("[PD5]")]}) g_tokenizer.add_special_tokens({"additional_special_tokens": [AddedToken("[YTMT]")]}) g_tokenizer.add_special_tokens({"additional_special_tokens": [AddedToken("[CLS]")]}) return g_tokenizer def _space_repl(m): s = m.group() n_newline = s.count("\n") if n_newline >= 2: return "\n\n" elif n_newline == 1: return "\n" return " " def _simplify_whitespace(text, retain_newlines=True): """simplify while respecting up to 2 newlines""" if retain_newlines: text = re.sub(r"\s+", _space_repl, text).strip() else: text = re.sub(r"\s+", " ", text).strip() return text def tokenize_batch( text_list, max_tokens=None, pad_token_id=0, retain_newlines=True, tokenizer_fp=None, ): tokenizer = _load_tokenizer(tokenizer_fp) text_list = [_simplify_whitespace(s, retain_newlines=retain_newlines) for s in text_list] text_enc = tokenizer( text_list, add_special_tokens=False, truncation=True, max_length=max_tokens, padding="longest", return_tensors="pt", )["input_ids"].type(torch.long) text_enc[text_enc == tokenizer.pad_token_id] = pad_token_id return text_enc def _clean_tag(tag): return re.sub(r"\s+", " ", tag).strip() def _augment_tag(s): if random.random() >= 0.95: s = s.upper() elif random.random() >= 0.95: s = s.capitalize() elif random.random() >= 0.9: s = s.title() elif random.random() >= 0.9: s = s.lower() if random.random() >= 0.5: s = s.replace("-", " ").strip() return s # Structure: # {start;vocals:start} # is song start & vocals start within 8s of actual start ## tags go here ## lyrics go here # {start;vocals:end} # is song end & vocals end within 8s of actual end def _get_start_control_tags(data_meta): start_s = data_meta.get("start_s") vocal_start_s = data_meta.get("vocal_start_s") control_tags = [] if start_s is not None and start_s <= 0.5: control_tags.append("start") if vocal_start_s is not None and vocal_start_s <= 8: control_tags.append("vocals:start") if len(control_tags) == 0: return None return "{" + ";".join(control_tags) + "}" def _get_end_control_tags(data_meta): end_s = data_meta.get("end_s") vocal_end_s = data_meta.get("vocal_end_s") original_duration_s = data_meta.get("original_duration_s") control_tags = [] if end_s is not None and original_duration_s - end_s <= 0.5: control_tags.append("end") if vocal_end_s is not None and end_s - vocal_end_s <= 10: control_tags.append("vocals:end") if len(control_tags) == 0: return None return "{" + ";".join(control_tags) + "}" def get_computed_tags(data_meta): computed_tags = [] cutoff_freq = data_meta.get("cutoff_freq") if cutoff_freq is None: return computed_tags if cutoff_freq <= 16_000: computed_tags.append("low rolloff") if cutoff_freq >= 18_000: computed_tags.append("high rolloff") return computed_tags def get_sample( data_sampling_info, split, dataset_idx=None, rel_row_idx=None, use_private=False, inference=False, suppress_text=False, dummy_data=False, return_idx=False, return_rel_row_idx=False, # for loading dpo data only ): if dummy_data: cfg = data_sampling_info["cfg"] x_audio_arr = np.zeros( ( cfg.semantic_n_codebooks + cfg.coarse_n_codebooks, cfg.block_size - cfg.t_text, ), dtype=np.int64, ) y_audio_arr = np.zeros( ( cfg.semantic_n_codebooks + cfg.coarse_n_codebooks, cfg.block_size - cfg.t_text, ), dtype=np.int64, ) return "", x_audio_arr, y_audio_arr data = data_sampling_info[split]["data"] metas = data_sampling_info[split]["metas"] idx_lists = data_sampling_info[split]["idx_lists"] if dataset_idx is None: weights = data_sampling_info[split]["weights"] dataset_idx = random.choices(list(range(len(weights))), weights=weights, k=1)[0] if rel_row_idx is None: rel_row_idx = random.choice(list(range(len(idx_lists[dataset_idx])))) # print("check", dataset_idx, rel_row_idx) rel_row_idx = rel_row_idx % len(idx_lists[dataset_idx]) row_idx = idx_lists[dataset_idx][rel_row_idx] # names = data_sampling_info[split]["names"] # dataset_name = names[dataset_idx] data_row = data[row_idx].astype(np.int64) data_meta = metas[row_idx] cfg = data_sampling_info["cfg"] # TODO: change the transpose here data_row = data_row.T # random mask is_mask = False if split == "train": if random.random() >= 0.3: is_mask = True num_mask = np.random.randint(1, int(cfg.t_audio * 0.4)) mask_indices = np.random.choice(cfg.t_audio, size=num_mask, replace=False) # build semantic y_semantic_arr = np.full( (cfg.semantic_n_codebooks, cfg.t_audio), cfg.semantic_pad_token, dtype=np.int64 ) for n in range(cfg.semantic_n_codebooks): y_semantic_arr[n, : data_row[n].shape[-1]] = data_row[n] if is_mask: y_semantic_arr[n, mask_indices] = cfg.semantic_pad_token # drop out semantic for better 'variations' model # if not inference and random.random() >= 0.75: # y_semantic_arr[:,:] = cfg.semantic_infer_token + 1 # build coarse y_coarse_arr = np.full((cfg.coarse_n_codebooks, cfg.t_audio), cfg.coarse_pad_token, dtype=np.int64) for n in range(cfg.coarse_n_codebooks): n2 = cfg.semantic_n_codebooks + n y_coarse_arr[n, : data_row[n2].shape[-1]] = data_row[n2] if is_mask: y_coarse_arr[n, mask_indices] = cfg.coarse_pad_token # combine audio and add x with infer token audio_arr = np.concatenate([y_semantic_arr, y_coarse_arr], axis=0) assert audio_arr.shape[-1] == cfg.block_size - cfg.t_text # build text text = "[BOS]" # if split == "train": # if random.random() >= 1.0: # always [BOS] # dataset_name = data_meta.get("dataset") # text = dataset_to_token[dataset_name] # collect tags if use_private: tags = data_meta.get("tags_private", data_meta.get("tags", [])) else: tags = data_meta.get("tags", []) # add computed tags computed_tags = get_computed_tags(data_meta) if len(computed_tags) > 0 and random.random() >= 0.1: tags.extend(computed_tags) # for tags remove newlines, empty tags, and case augment tags = [clean_tag for tag in tags if len(clean_tag := _clean_tag(tag)) > 0] long_tags = [tag.lower() for tag in tags if len(tag.split(" ")) > 3] # sentences tags = [tag.lower() for tag in tags if len(tag.split(" ")) < 3] # discard sentences tags = list(filter(lambda t: t not in BLACK_LIST, tags)) # remove irrelevant tags if (len(long_tags) > 0) and (random.random() >= 0.66): text = "[BOSL]" random.shuffle(long_tags) text += long_tags[0] elif len(tags) > 0: if inference: text += f"{', '.join(tags)}" else: random.shuffle(tags) tags = tags[:30] # max 30 tags tag_str = ", ".join(tags) text += tag_str # pretty arbitrary max len for now # get control tags text = text.replace("{", "").replace("}", "") text += "[EOS]" text = text.strip() if suppress_text: text = "" if return_idx: if return_rel_row_idx: return rel_row_idx, row_idx, text, x_audio_arr, y_audio_arr else: return row_idx, text, x_audio_arr, y_audio_arr return text, audio_arr def get_batch( data_sampling_info, split, dataset_idx=None, row_idx=None, use_private=False, inference=False, min_text_offs=None, suppress_text=False, dummy_data=False, return_idx=False, n_offs=None, load_dpo_pair=False, ): batch_size = data_sampling_info["batch_size"] device = data_sampling_info["device"] device_type = data_sampling_info["device_type"] tokenizer_fp = data_sampling_info.get("tokenizer_fp") cfg = data_sampling_info["cfg"] if not isinstance(dataset_idx, list): dataset_idx = [dataset_idx] * batch_size if not isinstance(row_idx, list): row_idx = [row_idx] * batch_size if n_offs is not None: row_idx = list(range(n_offs * batch_size, (n_offs + 1) * batch_size)) x_text_list = [] x_audio_list = [] y_list = [] idx_list = [] for n in range(batch_size): out = get_sample( data_sampling_info, split, dataset_idx=dataset_idx[n], rel_row_idx=row_idx[n], use_private=use_private, inference=inference, suppress_text=suppress_text, dummy_data=dummy_data, return_idx=return_idx, ) if return_idx: idx, x_text, x_audio, y = out idx_list.append(idx) else: x_text, x_audio = out x_text_list.append(x_text) x_audio_list.append(torch.from_numpy(x_audio)) x_text = tokenize_batch( x_text_list, max_tokens=cfg.t_text, pad_token_id=cfg.text_pad_token, tokenizer_fp=tokenizer_fp, ) # pad and concatenate x_audio = torch.stack(x_audio_list) # (batch, n_semantic + n_acoustic, n_audio) x = torch.concatenate( [ F.pad( x_text.unsqueeze(1), ( x_audio.shape[-1], cfg.block_size - x_audio.shape[-1] - x_text.shape[-1] + 1, ), "constant", cfg.text_pad_token, ), F.pad( x_audio[:, : cfg.semantic_n_codebooks], ( 0, cfg.block_size - x_audio.shape[-1] + 1, ), "constant", cfg.semantic_pad_token, ), F.pad( x_audio[:, cfg.semantic_n_codebooks :], ( 0, cfg.block_size - x_audio.shape[-1] + 1, ), "constant", cfg.coarse_pad_token, ), ], dim=1, ) audio_offset = x_audio.shape[-1] y = x.clone().detach()[:, :, audio_offset + 1 :] x = x[:, :, :-1] assert x.shape == ( batch_size, 1 + cfg.semantic_n_codebooks + cfg.coarse_n_codebooks, cfg.block_size, ) # crop only first six channels (text, semantic, 4 x acoustic) x = x[:, :6, :] y = y[:, :6, :] if device_type == "cuda": # pin arrays x,y, which allows us to move them to GPU asynchronously (non_blocking=True) x, y = ( x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True), ) else: x, y = x.to(device), y.to(device) del x_text_list, x_audio_list, y_list, x_text, x_audio return audio_offset, x, y