import random import re import numpy as np import torch import torch.nn.functional as F from data_utils_mmap import pad_x_arr, build_audio_arr, tokenize_batch, mask_batch_padding, build_text 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 abs_row_idx=None, ): if dummy_data: cfg = data_sampling_info["cfg"] x_audio_arr = np.zeros( ( cfg.semantic_n_codebooks + cfg.coarse_n_codebooks + 1, 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"] cfg = data_sampling_info["cfg"] 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 not None: # this is determined behavior rel_row_idx = rel_row_idx % len(idx_lists[dataset_idx]) row_idx = idx_lists[dataset_idx][rel_row_idx] else: row_idx = random.choice(idx_lists[dataset_idx]) rel_row_idx = idx_lists[dataset_idx].index(row_idx) # THIS IS A HARD OVERWRITE # GOD I HATE THIS if abs_row_idx is not None: row_idx = abs_row_idx def load_data_row(row_idx): data_row = data[row_idx].astype(np.int64) data_meta = metas[row_idx] # TODO: change the transpose here data_row = data_row.T # remove end padding if data_row.shape[0] == 1: # this is semantic only non_padding_indices = np.nonzero(data_row[0] != cfg.semantic_pad_token)[0] else: non_padding_indices = np.nonzero(data_row[-1] != cfg.coarse_pad_token)[0] last_non_padding_index = ( non_padding_indices[-1] if non_padding_indices.size > 0 else data_row.shape[-1] ) data_row = data_row[:, : last_non_padding_index + 1] return data_row, data_meta data_row, data_meta = load_data_row(row_idx) # disable control tags for DPO sample_tags = data_meta.get("tags", []) sample_text = data_meta.get("text", "") sample_neg_tags = data_meta.get("neg_tags", "") sample_gender_tags = data_meta.get("gender", "") sample_control_sliders = data_meta.get("control", {}) sample_control_tags = ( data_meta.get("control_tags", "") if isinstance(data_meta.get("control_tags", ""), str) else "" ) if len(sample_tags) > 0: current_tag = sample_tags[0] if sample_neg_tags: # we don't sample all of them for now for neg_tag in re.split(r"[,\.]", sample_neg_tags)[:5]: neg_tag = neg_tag.strip() if neg_tag: current_tag = current_tag + f", no {neg_tag}" if sample_gender_tags: current_tag = f"{sample_gender_tags}, {current_tag}" if sample_control_sliders and sample_control_tags == "": for slider_name in sorted(sample_control_sliders.keys()): slider_value = sample_control_sliders[slider_name] normalized_slider_value = int(round(slider_value * 10, 1)) sample_control_tags += f"{slider_name}:{normalized_slider_value};" sample_tags = [current_tag] if sample_text is None: print(f"WTF data_meta: {data_meta}") sample_text = "" if sample_tags is None: print(f"WTF data_meta: {data_meta}") sample_tags = "" if sample_control_tags: sample_control_tags = sample_control_tags.strip(";") sample_control_tags = "{" + sample_control_tags + "}" sample_duration_s = data_row.shape[-1] / cfg.semantic_rate_hz sample_vocal_start_s = data_meta.get("text_lines", [{}])[0].get("start_s", 0) text = build_text( sample_tags, sample_text, sample_duration_s, sample_vocal_start_s, inference, suppress_text, enable_control_tags=False, passin_control_tags=sample_control_tags, ) # if sample_neg_tags: # print(f"sample_neg_tags: {sample_neg_tags}") # print(f"text: {text}") x_text = tokenize_batch( [text], max_tokens=cfg.t_text, pad_token_id=cfg.text_pad_token, tokenizer_fp=data_sampling_info["tokenizer_fp"], )[0] # (N_TEXT_TOKENS) loss_start_index = 0 # if infilling -- data row is full concat arr # this is the infilling for auk # this is the prev configuration for 30b # if data_meta.get("task", None) == "infill": # suffix_len = data_meta["future_start_index"] # # .....future_start_index.... # # suffix is the future audio, prefix is the normal forward # suffix_audio_arr = build_audio_arr( # data_row[:, suffix_len:], cfg, semantic_infer_token=cfg.semantic_future_token # ) # prefix_audio_arr = build_audio_arr(data_row[:, :suffix_len], cfg) # x_audio_arr = np.concatenate([suffix_audio_arr, prefix_audio_arr], axis=-1) # # print(f"{suffix_audio_arr.shape=}, {prefix_audio_arr.shape=}, {x_audio_arr.shape=}") # # prepend future + generated index # loss_start_index = suffix_audio_arr.shape[-1] + data_meta["generated_start_index"] # cfg = data_sampling_info["cfg"] # if x_audio_arr.shape[-1] > cfg.t_audio: # raise ValueError(f"Sample overflows, {x_audio_arr.shape=}, {cfg.t_audio=}") # this is the configuration for auk if ( data_meta.get("task", None) == "infill" or data_meta.get("task", None) == "infill_intro" or data_meta.get("task", None) == "infill_outro" ): future_start_index = data_meta["future_start_index"] generation_start_index = data_meta["generated_start_index"] # .....future_start_index.... # suffix is the future audio, prefix is the normal forward suffix_audio_arr = build_audio_arr( data_row[:, future_start_index:], cfg, semantic_infer_token=cfg.semantic_future_token, include_eos=False, ) history_audio_arr = build_audio_arr( data_row[:, :generation_start_index], cfg, semantic_infer_token=cfg.semantic_history_token, include_eos=False, ) generated_audio_arr = build_audio_arr( data_row[:, generation_start_index:future_start_index], cfg, include_eos=True ) x_audio_arr = np.concatenate([suffix_audio_arr, history_audio_arr, generated_audio_arr], axis=-1) # print(f"{suffix_audio_arr.shape=}, {prefix_audio_arr.shape=}, {x_audio_arr.shape=}") # prepend future + generated index loss_start_index = suffix_audio_arr.shape[-1] + history_audio_arr.shape[-1] cfg = data_sampling_info["cfg"] if x_audio_arr.shape[-1] > cfg.t_audio: print(f"Sample overflows, {x_audio_arr.shape}, {cfg.t_audio}, data_meta: {data_meta}") x_audio_arr = x_audio_arr[:, : cfg.t_audio] elif data_meta.get("task", None) == "cover": generation_start_index = data_meta["generated_start_index"] generated_data_row = data_row[:, generation_start_index:] data_row_cover = data_row[:, :generation_start_index] cover_audio_arr = build_audio_arr( data_row_cover, cfg, semantic_infer_token=cfg.semantic_cover_token, include_eos=False, ) generated_audio_arr = build_audio_arr( generated_data_row, cfg, include_eos=True, ) loss_start_index = cover_audio_arr.shape[-1] # prepend cover to x_audio_arr x_audio_arr = np.concatenate( [cover_audio_arr, generated_audio_arr], axis=-1, ) elif data_meta.get("task", None) == "overpainting": generation_start_index = data_meta["generated_start_index"] generated_data_row = data_row[:, generation_start_index:] data_row_cover = data_row[:, :generation_start_index] overpaint_audio_arr = build_audio_arr( data_row_cover, cfg, semantic_infer_token=cfg.semantic_overpaint_token, include_eos=False, ) generated_audio_arr = build_audio_arr( generated_data_row, cfg, include_eos=True, ) loss_start_index = overpaint_audio_arr.shape[-1] # prepend cover to x_audio_arr x_audio_arr = np.concatenate( [overpaint_audio_arr, generated_audio_arr], axis=-1, ) elif data_meta.get("task", None) == "underpainting": generation_start_index = data_meta["generated_start_index"] generated_data_row = data_row[:, generation_start_index:] data_row_cover = data_row[:, :generation_start_index] underpaint_audio_arr = build_audio_arr( data_row_cover, cfg, semantic_infer_token=cfg.semantic_underpaint_token, include_eos=False, ) generated_audio_arr = build_audio_arr( generated_data_row, cfg, include_eos=True, ) loss_start_index = underpaint_audio_arr.shape[-1] # prepend cover to x_audio_arr x_audio_arr = np.concatenate( [underpaint_audio_arr, generated_audio_arr], axis=-1, ) elif data_meta.get("task", None) == "stem_condition": generation_start_index = data_meta["generated_start_index"] generated_data_row = data_row[:, generation_start_index:] data_row_stem = data_row[:, :generation_start_index] stem_audio_arr = build_audio_arr( data_row_stem, cfg, semantic_infer_token=cfg.semantic_stem_token, include_eos=False, ) generated_audio_arr = build_audio_arr( generated_data_row, cfg, include_eos=True, ) loss_start_index = stem_audio_arr.shape[-1] # prepend stem to x_audio_arr x_audio_arr = np.concatenate( [stem_audio_arr, generated_audio_arr], axis=-1, ) # Artist audio to audio elif data_meta.get("task", None) == "artist_consistency": generation_start_index = data_meta["generated_start_index"] generated_data_row = data_row[:, generation_start_index:] data_row_artist = data_row[:, :generation_start_index] artist_audio_arr = build_audio_arr( data_row_artist, cfg, semantic_infer_token=cfg.semantic_artist_token, include_eos=False, ) generated_audio_arr = build_audio_arr( generated_data_row, cfg, include_eos=True, ) loss_start_index = artist_audio_arr.shape[-1] x_audio_arr = np.concatenate( [ artist_audio_arr, generated_audio_arr, ], axis=-1, ) # Artist audio to audio elif data_meta.get("task", None) == "artist_cover": # the order is now: artist, cover, generated generation_start_index = data_meta["generated_start_index"] generated_data_row = data_row[:, generation_start_index:] cover_start_index = data_meta["cover_start_index"] data_row_artist = data_row[:, :cover_start_index] data_row_cover = data_row[:, cover_start_index:generation_start_index] cover_audio_arr = build_audio_arr( data_row_cover, cfg, semantic_infer_token=cfg.semantic_cover_token, include_eos=False ) artist_audio_arr = build_audio_arr( data_row_artist, cfg, semantic_infer_token=cfg.semantic_artist_token, include_eos=False ) generated_audio_arr = build_audio_arr( generated_data_row, cfg, include_eos=True, ) loss_start_index = artist_audio_arr.shape[-1] + cover_audio_arr.shape[-1] x_audio_arr = np.concatenate( [ artist_audio_arr, cover_audio_arr, generated_audio_arr, ], axis=-1, ) # do this for both extend and upload_extend elif data_meta.get("task", None) == "extend" or data_meta.get("task", None) == "upload_extend": # this time things are already ordered # however, we will only pendalize the loss from the right spot # give it 1 extra cause why not loss_start_index = max(data_meta["generated_start_index"] - 1, 0) x_audio_arr = build_audio_arr(data_row, cfg, include_eos=True) elif data_meta.get("task", None) == "cover_extend": # history_start_index is after the cover array # generated_start_index is after the history array history_start_index = data_meta["history_start_index"] data_row_cover = data_row[:, :history_start_index] generated_data_row = data_row[:, history_start_index:] cover_audio_arr = build_audio_arr( data_row_cover, cfg, semantic_infer_token=cfg.semantic_cover_token, include_eos=False ) generated_audio_arr = build_audio_arr(generated_data_row, cfg, include_eos=True) x_audio_arr = np.concatenate( [ cover_audio_arr, generated_audio_arr, ], axis=-1, ) loss_start_index = cover_audio_arr.shape[-1] + max(data_meta["generated_start_index"], 0) # history_start_index is after the artist array # generated_start_index is after the history array elif data_meta.get("task", None) == "artist_extend": history_start_index = data_meta["history_start_index"] data_row_artist = data_row[:, :history_start_index] generated_data_row = data_row[:, history_start_index:] artist_audio_arr = build_audio_arr( data_row_artist, cfg, semantic_infer_token=cfg.semantic_artist_token, include_eos=False ) generated_audio_arr = build_audio_arr(generated_data_row, cfg, include_eos=True) x_audio_arr = np.concatenate( [ artist_audio_arr, generated_audio_arr, ], axis=-1, ) loss_start_index = artist_audio_arr.shape[-1] + max(data_meta["generated_start_index"], 0) elif data_meta.get("task", None) == "playlist_condition": generation_start_index = data_meta["generated_start_index"] generated_data_row = data_row[:, generation_start_index:] data_row_playlist = data_row[:, :generation_start_index] playlist_audio_arr = [] start_idx = 0 for array_len in data_meta["playlist_arr_len"]: if array_len == 0: continue # end_idx = min(start_idx + array_len, data_row_playlist.shape[-1]) playlist_audio_arr.append( build_audio_arr( data_row_playlist[:, start_idx:end_idx], cfg, semantic_infer_token=cfg.semantic_playlist_token, include_eos=False, ) ) start_idx = end_idx # Increment by actual amount used generated_audio_arr = build_audio_arr( generated_data_row, cfg, include_eos=True, ) concat_playlist_audio_arr = np.concatenate(playlist_audio_arr, axis=-1) loss_start_index = concat_playlist_audio_arr.shape[-1] x_audio_arr = np.concatenate( [ concat_playlist_audio_arr, generated_audio_arr, ], axis=-1, ) else: loss_start_index = 0 # this is for normal generation x_audio_arr = build_audio_arr(data_row, cfg, include_eos=True) # construct full x arr x_arr = np.empty( ( 1 + cfg.semantic_n_codebooks + cfg.coarse_n_codebooks, len(x_text) + x_audio_arr.shape[-1], ), dtype=np.int64, ) x_arr[0, : len(x_text)] = x_text x_arr[0, len(x_text) :] = cfg.text_pad_token x_arr[1 : 1 + cfg.semantic_n_codebooks, : len(x_text)] = cfg.semantic_pad_token x_arr[1 + cfg.semantic_n_codebooks :, : len(x_text)] = cfg.coarse_pad_token x_arr[1:, len(x_text) :] = x_audio_arr # we should maybe plus one ... but there is shift 1 already loss_start_index += len(x_text) if return_idx: if return_rel_row_idx: return rel_row_idx, row_idx, text, x_arr, loss_start_index else: return row_idx, text, x_arr, loss_start_index return text, x_arr, loss_start_index 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, abs_row_idx=None, ): 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)) if not isinstance(abs_row_idx, list): abs_row_idx = [abs_row_idx] * batch_size x_text_list = [] x_list = [] y_list = [] idx_list = [] loss_start_index_list = [] if not load_dpo_pair: 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, abs_row_idx=abs_row_idx[n], ) if return_idx: idx, x_text, x_arr, x_loss_start_index = out idx_list.append(idx) else: x_text, x_arr, x_loss_start_index = out x_text_list.append(x_text) # truncate to block size for now still, even with cover, artist, etc x_arr = x_arr[:, : cfg.block_size] # (C, T) x_arr = pad_x_arr(x_arr, cfg) x_list.append(x_arr) loss_start_index_list.append(x_loss_start_index) else: assert batch_size % 2 == 0 for n in range(batch_size // 2): out = get_sample( data_sampling_info, split, dataset_idx=0, # negative is always the first dataset rel_row_idx=row_idx[n], use_private=use_private, inference=True, # DPO is inference! without text augmentation suppress_text=suppress_text, dummy_data=dummy_data, return_idx=True, return_rel_row_idx=True, # use this to fetch the positive sample ) neg_rel_row_idx, neg_row_idx, x_text, x_arr, x_loss_start_index = out # print(row_idx[n], x_text) x_text_list.append(x_text) # truncate to block size for now still, even with cover, artist, etc x_arr = x_arr[:, : cfg.block_size] # (C, T) x_arr = pad_x_arr(x_arr, cfg) x_list.append(x_arr) idx_list.append(neg_row_idx) loss_start_index_list.append(x_loss_start_index) # now load positive out = get_sample( data_sampling_info, split, dataset_idx=1, # positive is always the 2nd dataset rel_row_idx=neg_rel_row_idx, use_private=use_private, inference=True, suppress_text=suppress_text, dummy_data=dummy_data, return_idx=True, return_rel_row_idx=True, ) pos_rel_row_idx, pos_row_idx, pos_x_text, pos_x_arr, pos_x_loss_start_index = out x_text_list.append(pos_x_text) # truncate to block size for now still, even with cover, artist, etc pos_x_arr = pos_x_arr[:, : cfg.block_size] # (C, T) pos_x_arr = pad_x_arr(pos_x_arr, cfg) if pos_x_text != x_text: print( "WARNING: prompt mismatch in DPO pair", neg_rel_row_idx, pos_rel_row_idx, neg_row_idx, pos_row_idx, pos_x_text, x_text, ) if pos_x_loss_start_index != x_loss_start_index: print( "WARNING: loss start index mismatch in DPO pair", neg_rel_row_idx, pos_rel_row_idx, neg_row_idx, pos_row_idx, pos_x_loss_start_index, x_loss_start_index, ) x_list.append(pos_x_arr) idx_list.append(pos_row_idx) loss_start_index_list.append(pos_x_loss_start_index) x = np.stack(x_list, axis=0) assert x.shape == ( batch_size, 1 + cfg.semantic_n_codebooks + cfg.coarse_n_codebooks, cfg.block_size, ), x.shape y = x.copy()[:, 1:, 1:] # remove text stream and first token x = torch.from_numpy(x) mask_batch_padding(y, cfg) y = torch.from_numpy(y) 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_list, y_list, x_text if return_idx: return idx_list, x, y, loss_start_index_list return x, y, loss_start_index_list # copied from https://github.com/eric-mitchell/direct-preference-optimization/blob/main/trainers.py # TODO: need a lot more work on this...X.x totally misunderstood # KTO: https://github.com/ContextualAI/HALOs/blob/main/trainers.py#L742 def preference_loss( policy_chosen_logps: torch.FloatTensor, policy_rejected_logps: torch.FloatTensor, reference_chosen_logps: torch.FloatTensor, reference_rejected_logps: torch.FloatTensor, beta: float, label_smoothing: float = 0.0, ipo: bool = False, reference_free: bool = False, kto: bool = False, ): """Compute the DPO loss for a batch of policy and reference model log probabilities. Args: policy_chosen_logps: Log probabilities of the policy model for the chosen responses. Shape: (batch_size,) policy_rejected_logps: Log probabilities of the policy model for the rejected responses. Shape: (batch_size,) reference_chosen_logps: Log probabilities of the reference model for the chosen responses. Shape: (batch_size,) reference_rejected_logps: Log probabilities of the reference model for the rejected responses. Shape: (batch_size,) beta: Temperature parameter for the DPO loss, typically something in the range of 0.1 to 0.5. We ignore the reference model as beta -> 0. label_smoothing: conservativeness for DPO loss, which assumes that preferences are noisy (flipped with probability label_smoothing) ipo: If True, use the IPO loss instead of the DPO loss. reference_free: If True, we ignore the _provided_ reference model and implicitly use a reference model that assigns equal probability to all responses. Returns: A tuple of three tensors: (losses, chosen_rewards, rejected_rewards). The losses tensor contains the DPO loss for each example in the batch. The chosen_rewards and rejected_rewards tensors contain the rewards for the chosen and rejected responses, respectively. """ if not kto: pi_logratios = policy_chosen_logps - policy_rejected_logps ref_logratios = reference_chosen_logps - reference_rejected_logps if reference_free: ref_logratios = 0 logits = pi_logratios - ref_logratios # also known as h_{\pi_\theta}^{y_w,y_l} if ipo: losses = (logits - 1 / (2 * beta)) ** 2 # Eq. 17 of https://arxiv.org/pdf/2310.12036v2.pdf else: # Eq. 3 https://ericmitchell.ai/cdpo.pdf; label_smoothing=0 gives original DPO (Eq. 7 of https://arxiv.org/pdf/2305.18290.pdf) losses = ( -F.logsigmoid(beta * logits) * (1 - label_smoothing) - F.logsigmoid(-beta * logits) * label_smoothing ) elif kto: chosen_logratios = policy_chosen_logps - reference_chosen_logps rejected_logratios = policy_rejected_logps - reference_rejected_logps chosen_KL = chosen_logratios.mean().clamp(min=0) rejected_KL = rejected_logratios.mean().clamp(min=0) # simplified from the original KTO loss losses = -F.sigmoid(beta * (chosen_logratios - rejected_KL)) - F.sigmoid( beta * (chosen_KL - rejected_logratios) ) # these are the same chosen_rewards = beta * (policy_chosen_logps - reference_chosen_logps).detach() rejected_rewards = beta * (policy_rejected_logps - reference_rejected_logps).detach() return losses, chosen_rewards, rejected_rewards