1410 lines
79 KiB
Python
1410 lines
79 KiB
Python
import random
|
|
|
|
import numpy as np
|
|
import logging
|
|
import argparse, copy
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from torchmetrics.classification import MulticlassAccuracy
|
|
|
|
from .modules.utils import make_pad_mask
|
|
|
|
from .modules.embedding import SinePositionalEmbedding, TokenEmbedding
|
|
from .modules.transformer import (
|
|
LayerNorm,
|
|
TransformerEncoder,
|
|
TransformerEncoderLayer,
|
|
)
|
|
from .codebooks_patterns import DelayedPatternProvider
|
|
|
|
def top_k_top_p_filtering(
|
|
logits, top_k=0, top_p=1.0, filter_value=-float("Inf"), min_tokens_to_keep=1
|
|
):
|
|
"""Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
|
|
Args:
|
|
logits: logits distribution shape (batch size, vocabulary size)
|
|
if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
|
|
if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
|
|
Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
|
|
Make sure we keep at least min_tokens_to_keep per batch example in the output
|
|
From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
|
|
"""
|
|
if top_k > 0:
|
|
top_k = min(
|
|
max(top_k, min_tokens_to_keep), logits.size(-1)
|
|
) # Safety check
|
|
# Remove all tokens with a probability less than the last token of the top-k
|
|
indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
|
|
logits[indices_to_remove] = filter_value
|
|
|
|
if top_p < 1.0:
|
|
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
|
|
cumulative_probs = torch.cumsum(
|
|
F.softmax(sorted_logits, dim=-1), dim=-1
|
|
)
|
|
|
|
# Remove tokens with cumulative probability above the threshold (token with 0 are kept)
|
|
sorted_indices_to_remove = cumulative_probs > top_p
|
|
if min_tokens_to_keep > 1:
|
|
# Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
|
|
sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
|
|
# Shift the indices to the right to keep also the first token above the threshold
|
|
sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[
|
|
..., :-1
|
|
].clone()
|
|
sorted_indices_to_remove[..., 0] = 0
|
|
|
|
# scatter sorted tensors to original indexing
|
|
indices_to_remove = sorted_indices_to_remove.scatter(
|
|
1, sorted_indices, sorted_indices_to_remove
|
|
)
|
|
logits[indices_to_remove] = filter_value
|
|
return logits
|
|
|
|
|
|
def topk_sampling(logits, top_k=10, top_p=1.0, temperature=1.0):
|
|
# temperature: (`optional`) float
|
|
# The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
|
|
# top_k: (`optional`) int
|
|
# The number of highest probability vocabulary tokens to keep for top-k-filtering. Between 1 and infinity. Default to 50.
|
|
# top_p: (`optional`) float
|
|
# The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Must be between 0 and 1. Default to 1.
|
|
|
|
# Temperature (higher temperature => more likely to sample low probability tokens)
|
|
if temperature != 1.0:
|
|
logits = logits / temperature
|
|
# Top-p/top-k filtering
|
|
logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
|
|
# Sample
|
|
token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
|
|
return token
|
|
|
|
|
|
|
|
class VoiceCraft(nn.Module):
|
|
def __init__(self, args):
|
|
super().__init__()
|
|
self.args = copy.copy(args)
|
|
self.pattern = DelayedPatternProvider(n_q=self.args.n_codebooks)
|
|
if not getattr(self.args, "special_first", False):
|
|
self.args.special_first = 0
|
|
if not getattr(self.args, "n_special", False):
|
|
self.args.n_special = 3
|
|
self.args.eos = getattr(self.args, "eos", -1)
|
|
self.eog = nn.Parameter(torch.full((self.args.n_codebooks, 1), self.args.eog, dtype=torch.long), requires_grad=False) # [K 1]
|
|
if self.args.eos > 0:
|
|
assert self.args.eos != self.args.audio_pad_token and self.args.eos != self.args.empty_token, self.args.eos
|
|
self.eos = nn.Parameter(torch.full((self.args.n_codebooks, 1), self.args.eos, dtype=torch.long), requires_grad=False) # [K 1]
|
|
if type(self.args.audio_vocab_size) == str:
|
|
self.args.audio_vocab_size = eval(self.args.audio_vocab_size)
|
|
|
|
self.n_text_tokens = self.args.text_vocab_size + 1
|
|
assert self.args.text_pad_token == self.args.text_vocab_size, f"self.args.text_vocab_size: {self.args.text_vocab_size}, self.args.text_pad_token: {self.args.text_pad_token}"
|
|
|
|
self.n_audio_tokens = [self.args.audio_vocab_size + self.args.n_special] * self.args.n_codebooks # special tokens: empty token, EOG token, audio pad token
|
|
assert self.args.audio_vocab_size == self.args.empty_token, self.args.empty_token
|
|
assert self.args.eog == self.args.audio_vocab_size + 1, self.args.eog
|
|
assert self.args.audio_pad_token == self.args.audio_vocab_size + 2, self.args.audio_pad_token
|
|
|
|
self.text_embedding = TokenEmbedding(
|
|
dim_model=self.args.d_model,
|
|
vocab_size=self.n_text_tokens,
|
|
dropout=self.args.text_embedding_dropout
|
|
)
|
|
|
|
self.audio_embedding = nn.ModuleList(
|
|
[
|
|
TokenEmbedding(
|
|
dim_model=self.args.audio_embedding_dim,
|
|
vocab_size=self.n_audio_tokens[k],
|
|
dropout=self.args.audio_embedding_dropout
|
|
) for k in range(self.args.n_codebooks)
|
|
]
|
|
)
|
|
self.mask_embedding = nn.Parameter(torch.randn(self.args.max_n_spans, self.args.d_model), requires_grad=True)
|
|
self.text_positional_embedding = SinePositionalEmbedding(
|
|
self.args.d_model,
|
|
dropout=self.args.text_positional_embedding_dropout,
|
|
scale=False,
|
|
alpha=True, # learnable scaler, scale the volume of positional embedding
|
|
)
|
|
self.audio_positional_embedding = SinePositionalEmbedding(
|
|
self.args.d_model,
|
|
dropout=self.args.audio_positional_embedding_dropout,
|
|
scale=False,
|
|
alpha=True, # learnable scaler, scale the volume of positional embedding
|
|
)
|
|
|
|
dec_layer = TransformerEncoderLayer(
|
|
self.args.d_model,
|
|
self.args.nhead,
|
|
dim_feedforward=self.args.d_model * 4,
|
|
dropout=self.args.trm_dropout,
|
|
batch_first=True,
|
|
norm_first=True,
|
|
layer_norm_cls=LayerNorm
|
|
)
|
|
self.decoder = TransformerEncoder(
|
|
dec_layer,
|
|
num_layers=self.args.num_decoder_layers,
|
|
norm=LayerNorm(self.args.d_model),
|
|
)
|
|
|
|
self.predict_layer = nn.ModuleList(
|
|
[
|
|
nn.Sequential(nn.Linear(self.args.d_model, self.args.audio_vocab_size//2), nn.GELU(), nn.Linear(self.args.audio_vocab_size//2, self.n_audio_tokens[k])) for k in range(self.args.n_codebooks)
|
|
]
|
|
)
|
|
|
|
self.accuracy_metrics = nn.ModuleList(
|
|
[MulticlassAccuracy(
|
|
self.n_audio_tokens[k],
|
|
top_k=10,
|
|
average="micro",
|
|
multidim_average="global",
|
|
ignore_index=None,
|
|
) for k in range(self.args.n_codebooks)]
|
|
)
|
|
|
|
|
|
def prepare_mask_intervals(self, y_lens):
|
|
mask_intervals = []
|
|
non_mask_intervals = []
|
|
|
|
for i, y_len in enumerate(y_lens):
|
|
if self.args.mask_sample_dist == "uniform":
|
|
n_spans = random.choice(range(1, self.args.max_n_spans+1))
|
|
elif "poisson" in self.args.mask_sample_dist.lower():
|
|
param = float(self.args.mask_sample_dist[len("poisson"):])
|
|
poisson_sample = torch.poisson(torch.tensor([param]))
|
|
n_spans = int(poisson_sample.clamp(1, self.args.max_n_spans).item())
|
|
|
|
starts = random.sample(range(1, y_len-1-self.args.mask_len_min), n_spans)
|
|
starts = sorted(starts)
|
|
|
|
for j in range(len(starts)-1, 0, -1):
|
|
if starts[j] - starts[j-1] < self.args.min_gap:
|
|
del starts[j] # If elements are too close, delete the later one
|
|
assert len(starts) > 0, f"there is no masked span left, y_len: {y_len}, sampled n_spans: {n_spans}"
|
|
|
|
temp_starts = starts + [y_len]
|
|
gaps = [temp_starts[j+1] - temp_starts[j] for j in range(len(temp_starts)-1)]
|
|
|
|
ends = []
|
|
|
|
for j, (start, gap) in enumerate(zip(starts, gaps)):
|
|
mask_len = random.randint(self.args.mask_len_min, self.args.mask_len_max)
|
|
# if mask_len > gap * self.args.max_mask_portion: # make sure the masks are not overlapping with each other
|
|
if mask_len > gap - 1: # make sure the masks are not overlapping with each other
|
|
# temp_mask_start = int(0.6*gap*self.args.max_mask_portion)
|
|
# temp_mask_end = int(gap*self.args.max_mask_portion)
|
|
temp_mask_start = 1
|
|
temp_mask_end = gap - 1
|
|
mask_len = random.randint(temp_mask_start, temp_mask_end)
|
|
ends.append(start + mask_len)
|
|
|
|
mask_intervals.append([(s,e) for s,e in zip(starts, ends)])
|
|
non_mask_intervals.append([(ns,ne) for ns, ne in zip([0]+ends, starts+[y_len])])
|
|
|
|
return mask_intervals, non_mask_intervals
|
|
|
|
def rearrange(self, y, non_mask_intervals, mask_intervals):
|
|
reduced_eog = getattr(self.args, "reduced_eog", 0)
|
|
rearranged_y = []
|
|
for i in range(len(y)):
|
|
if self.args.eos > 0:
|
|
assert reduced_eog
|
|
cur_y = [y[i, :, item[0]: item[1]] for item in non_mask_intervals[i][:-1]] + [torch.cat([y[i, :, non_mask_intervals[i][-1][0]: non_mask_intervals[i][-1][1]], self.eos], dim=-1)] + [torch.cat([y[i, :, item[0]: item[1]], self.eog], dim=-1) for item in mask_intervals[i]] # only insert eog to the last non-mask-interval, which is when the utterance actual ends
|
|
else:
|
|
if reduced_eog:
|
|
cur_y = [y[i, :, item[0]: item[1]] for item in non_mask_intervals[i][:-1]] + [torch.cat([y[i, :, non_mask_intervals[i][-1][0]: non_mask_intervals[i][-1][1]], self.eog], dim=-1)] + [torch.cat([y[i, :, item[0]: item[1]], self.eog], dim=-1) for item in mask_intervals[i]] # only insert eog to the last non-mask-interval, which is when the utterance actual ends
|
|
else:
|
|
cur_y = [torch.cat([y[i, :, item[0]: item[1]], self.eog], dim=-1) for item in non_mask_intervals[i]] + [torch.cat([y[i, :, item[0]: item[1]], self.eog], dim=-1) for item in mask_intervals[i]] # eog is added to each section TODO this is not correct, I should add eog to non_mask_intervals if that segment is not the ending segment (as there is no way for the model to predict eog for those segments, and this will do harm to tts experiment, where the model randomly output eog for the first segment)
|
|
rearranged_y.append(cur_y)
|
|
return rearranged_y
|
|
|
|
def shift(self, rearranged_y):
|
|
shifted_y = []
|
|
patterns = []
|
|
for i in range(len(rearranged_y)):
|
|
cur_patterns = [self.pattern.get_pattern(cur_y.shape[1]) for cur_y in rearranged_y[i]]
|
|
out = [cur_pattern.build_pattern_sequence(z=cur_y.unsqueeze(0).contiguous(), special_token=self.args.empty_token, keep_only_valid_steps=False) for cur_pattern, cur_y in zip(cur_patterns, rearranged_y[i])]
|
|
shifted_y.append([item[0].squeeze(0) for item in out]) # the first item is values, later two are indexes and mask
|
|
patterns.append(cur_patterns)
|
|
return shifted_y, patterns
|
|
|
|
def insert_mask(self, shifted_y):
|
|
inserted_y = []
|
|
mask_position = []
|
|
mask_value = []
|
|
for i in range(len(shifted_y)):
|
|
num_masks = (len(shifted_y[i]) - 1) // 2
|
|
assert num_masks == (len(shifted_y[i]) - 1) / 2, len(shifted_y[i])
|
|
emb_inds = list(range(self.args.max_n_spans))
|
|
if self.args.shuffle_mask_embedding:
|
|
random.shuffle(emb_inds)
|
|
emb_inds_use = emb_inds[:num_masks]
|
|
emb_inds_use = emb_inds_use + emb_inds_use
|
|
mask_value.append(emb_inds_use)
|
|
cur_inserted_y = []
|
|
cur_mask_position = []
|
|
for j in range(len(shifted_y[i])-1):
|
|
cur_inserted_y.append(shifted_y[i][j])
|
|
cur_mask_position.append(sum([item.shape[1] for item in cur_inserted_y])) # each item is of shape [K S], so take shape[1]
|
|
cur_inserted_y.append(self.eog) # insert mask token of shape [K, 1], BUT we are actually using the eog token as a place holder here, as the real mask will be inserted in embed_y function
|
|
|
|
cur_inserted_y.append(shifted_y[i][-1])
|
|
|
|
inserted_y.append(cur_inserted_y)
|
|
mask_position.append(cur_mask_position)
|
|
return inserted_y, mask_position, mask_value
|
|
|
|
def cat_y(self, inserted_y, mask_position, y_lens):
|
|
reduced_eog = getattr(self.args, "reduced_eog", 0)
|
|
cated_y = []
|
|
new_y_lens = []
|
|
for i in range(len(inserted_y)):
|
|
cur_cated_y = torch.cat(inserted_y[i], dim=1) #[K S]
|
|
cur_cated_y = cur_cated_y.transpose(1,0) # [S K]
|
|
cur_cated_y_len = cur_cated_y.shape[0]
|
|
if reduced_eog:
|
|
assert cur_cated_y_len == y_lens[i] + len(mask_position[i]) + (len(mask_position[i]) + 1) * self.args.n_codebooks + (len(mask_position[i])/2 + 1), f"cur_cated_y_len == {cur_cated_y_len}, but it should be y_lens[i] ({y_lens[i]}) + len(mask_position[i]) ({len(mask_position[i])}) + (len(mask_position[i]) + 1) * self.args.n_codebooks ({(len(mask_position[i]) + 1) * self.args.n_codebooks}) + (len(mask_position[i])/2 + 1) ({len(mask_position[i])/2 + 1})={y_lens[i] + len(mask_position[i]) + (len(mask_position[i]) + 1) * self.args.n_codebooks + (len(mask_position[i])/2 + 1)}"
|
|
else:
|
|
assert cur_cated_y_len == y_lens[i] + len(mask_position[i]) + (len(mask_position[i]) + 1) * self.args.n_codebooks + (len(mask_position[i]) + 1), f"cur_cated_y_len == {cur_cated_y_len}, but it should be y_lens[i] ({y_lens[i]}) + len(mask_position[i]) ({len(mask_position[i])}) + (len(mask_position[i]) + 1) * self.args.n_codebooks ({(len(mask_position[i]) + 1) * self.args.n_codebooks}) + (len(mask_position[i]) + 1) ({len(mask_position[i]) + 1})" # the last term represent the inserted eog token, originally it's inserted at the end of every token, but this is wrong
|
|
new_y_lens.append(cur_cated_y_len)
|
|
cated_y.append(cur_cated_y)
|
|
|
|
cated_y = torch.nn.utils.rnn.pad_sequence(cated_y, batch_first=False, padding_value=self.args.audio_pad_token)
|
|
assert cated_y.shape == torch.Size([max(new_y_lens),len(inserted_y), self.args.n_codebooks]), f"cated_y.shape: {cated_y.shape}, but it should be {torch.Size([max(new_y_lens,len(inserted_y), self.args.n_codebooks)])}"
|
|
cated_y = cated_y.permute(2,0,1) # [T,B,K]->[K,T,B]
|
|
assert cated_y.shape[0] == self.args.n_codebooks, cated_y.shape
|
|
return cated_y, torch.LongTensor(new_y_lens).to(cated_y.device)
|
|
|
|
def embed_y(self, cated_y, mask_position, mask_value):
|
|
embedded_y = torch.stack([self.audio_embedding[k](cated_y[k]) for k in range(self.args.n_codebooks)], dim=0) # [K, T, B, D]
|
|
assert embedded_y.shape[0] == self.args.n_codebooks, embedded_y.shape
|
|
assert embedded_y.shape[-1] == self.args.d_model, embedded_y.shape
|
|
embedded_y = embedded_y.sum(dim=0) # [K,T,B,D]->[T,B,D]
|
|
embedded_y = embedded_y.transpose(1,0) # [T,B,D]->[B,T,D]
|
|
for i in range(len(embedded_y)):
|
|
if len(mask_position[i]) > 0:
|
|
embedded_y[i, mask_position[i]] = self.mask_embedding[mask_value[i]]
|
|
return embedded_y
|
|
|
|
def prepare_input_target(self, y, y_lens):
|
|
# rearrange y
|
|
# assume y shape: [B T K], K is n_codebooks
|
|
assert y.shape[1] == self.args.n_codebooks, y.shape
|
|
# sample mask_intervals
|
|
mask_intervals, non_mask_intervals = self.prepare_mask_intervals(y_lens)
|
|
|
|
# need to have EOG in each section (SOG will be generated by the pattern class)
|
|
# but mask can be inserted later after we have shifted the input
|
|
# y could be rearranged in this way:
|
|
# [
|
|
# [tensor[4, 12], tensor[4, 45], tensor[4, 102], tensor[4, 32]], tensor[4, 22]],
|
|
# [tensor[4, 44], tensor[4, 56], tensor[4, 19]],
|
|
# ...
|
|
# ]
|
|
# for the first list of tensors (4 tensors), first 3 tensors are non_masked part, last 2 are masked part.
|
|
# NOTE #non_masked_part = #masked_part + 1
|
|
# NOTE *these are also the targets*
|
|
# added eog at the end of each segment (masked segment and unmasked segment)
|
|
rearranged_y = self.rearrange(y, non_mask_intervals, mask_intervals)
|
|
targets = rearranged_y # each element in each sample is of shape [K T]
|
|
assert targets[0][0].shape[0] == self.args.n_codebooks, targets[0][0].shape
|
|
|
|
# next we need to apply pattern shifting to each tensor, after which, we'll replace the starting tokens of each section with a token that's different from the special padding token
|
|
# [[5, 1, 2, 3, 4, 5, 5],
|
|
# [5, 5, 1, 2, 3, 4, 5],
|
|
# [5, 5, 5, 1, 2, 3, 4]]
|
|
shifted_y, patterns = self.shift(rearranged_y) # each element [K S]
|
|
assert shifted_y[0][0].shape[0] == self.args.n_codebooks, shifted_y[0][0].shape[0]
|
|
|
|
|
|
# then, insert mask token at the intersection of each tensor (we want to decide the arrangement of the mask (shuffle or not)), we better have a separate nn.embedding for it
|
|
# we also need to record the position of the inserted mask
|
|
inserted_y, mask_position, mask_value = self.insert_mask(shifted_y)
|
|
assert inserted_y[0][0].shape[0] == self.args.n_codebooks, inserted_y[0][0].shape[0]
|
|
assert inserted_y[0][1].shape == torch.Size((self.args.n_codebooks, 1)), f"this should be a mask, so should have shape {(self.args.n_codebooks, 1)}, but it's {inserted_y[0][1].shape}"
|
|
|
|
# then concat tensors that belong to the same sample (in order) then get the length of each sample, and then stack them in batch dimension, pad them with pad_token
|
|
cated_y, new_y_lens = self.cat_y(inserted_y, mask_position, y_lens) # KTB
|
|
assert cated_y.shape == torch.Size((self.args.n_codebooks, cated_y.shape[1], len(inserted_y)))
|
|
|
|
|
|
# embed remember to separately embed the mask tokens
|
|
embedded_y = self.embed_y(cated_y, mask_position, mask_value) #BTD
|
|
assert embedded_y.shape[1:] == torch.Size((max(new_y_lens), self.args.d_model)), embedded_y.shape
|
|
|
|
# positional embedding
|
|
y_input = self.audio_positional_embedding(embedded_y)
|
|
|
|
# make attention mask and padding mask
|
|
y_padding_mask = make_pad_mask(new_y_lens).to(y.device)
|
|
y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y_padding_mask.device)
|
|
return y_input, new_y_lens, targets, y_padding_mask, y_attention_mask, mask_position, patterns
|
|
|
|
def remove_mask(self, logits, mask_position, new_y_lens):
|
|
# logits: [B K S card]
|
|
logits_use = []
|
|
for i in range(len(logits)):
|
|
non_mask_positions = [-1] + mask_position[i] + [new_y_lens[i]]
|
|
non_mask_intervals = [[non_mask_positions[i]+1, non_mask_positions[i+1]] for i in range(len(non_mask_positions)-1)]
|
|
cur_logits_use = [logits[i, :, l:r] for l,r in non_mask_intervals]
|
|
logits_use.append(cur_logits_use)
|
|
|
|
return logits_use
|
|
|
|
def revert_pattern(self, patterns, logits_use):
|
|
logits_final = []
|
|
logit_masks = []
|
|
for i in range(len(logits_use)):
|
|
cur_logits = [
|
|
item.unsqueeze(0).permute(0, 3, 1, 2).contiguous() for item in logits_use[i]
|
|
] # each item is of shape [1 K S card] [1 card K S]
|
|
cur_logits_final = [
|
|
cur_pattern.revert_pattern_logits(
|
|
item, 0, keep_only_valid_steps=False
|
|
)
|
|
for cur_pattern, item in zip(patterns[i], cur_logits)
|
|
] # if input output order doesn't match, this step will give an error
|
|
cur_logits_final_ret = [item[0].permute(0,2,3,1).squeeze(0) for item in cur_logits_final] # each element is of shape [K,T,card]
|
|
logits_final.append(cur_logits_final_ret)
|
|
logit_masks.append([item[2] for item in cur_logits_final])
|
|
|
|
return logits_final, logit_masks
|
|
|
|
def dec_forward(
|
|
self,
|
|
x_input,
|
|
x_lens,
|
|
x_attention_mask,
|
|
x_padding_mask,
|
|
y_input,
|
|
new_y_lens,
|
|
y_attention_mask,
|
|
y_padding_mask,
|
|
past=None,
|
|
last_3_tokens=False
|
|
):
|
|
x_attn_mask = F.pad(
|
|
x_attention_mask,
|
|
(0, new_y_lens.max()),
|
|
value=True,
|
|
) # x attn to all x, doesn't attn to any y, this follow figure 3 of the valle paper
|
|
y_attn_mask = F.pad(
|
|
y_attention_mask,
|
|
(x_lens.max(), 0), # y is padded at the front
|
|
value=False,
|
|
) # y attn to all x, for y itself use lower triangle mask to ensure autoregressive
|
|
xy_attn_mask = torch.concat([x_attn_mask, y_attn_mask], dim=0)
|
|
|
|
# merge key padding and attention masks
|
|
bsz, src_len = x_input.shape[0], x_lens.max() + new_y_lens.max()
|
|
xy_padding_mask = torch.concat([x_padding_mask, y_padding_mask], dim=1)
|
|
_xy_padding_mask = (
|
|
xy_padding_mask.view(bsz, 1, 1, src_len)
|
|
.expand(-1, self.args.nhead, -1, -1)
|
|
.reshape(bsz * self.args.nhead, 1, src_len)
|
|
)
|
|
# Check shapes and resize or broadcast as necessary
|
|
if xy_attn_mask.shape != _xy_padding_mask.shape:
|
|
# Assuming _xy_padding_mask has the correct shape and xy_attn_mask needs adjustment
|
|
xy_attn_mask = xy_attn_mask.expand_as(_xy_padding_mask) # Example approach
|
|
xy_attn_mask = xy_attn_mask.logical_or(_xy_padding_mask)
|
|
|
|
new_attn_mask = torch.zeros_like(xy_attn_mask)
|
|
new_attn_mask.masked_fill_(xy_attn_mask, float("-inf"))
|
|
xy_attn_mask = new_attn_mask
|
|
|
|
xy_input = torch.cat([x_input, y_input], dim=1)
|
|
|
|
if past == None: # do not use kvcache
|
|
out, _ = self.decoder((xy_input, None), mask=xy_attn_mask)
|
|
return out[:, x_lens.max():], None
|
|
else: # use kvcache
|
|
if past.ndim > 3: # uses kvcache, only need to pass the last tokens, this doesn't work with multi-span speech editing yet
|
|
if last_3_tokens:
|
|
xy_input = xy_input[:, -3:]
|
|
xy_attn_mask = xy_attn_mask[:, -3:]
|
|
else:
|
|
xy_input = xy_input[:, -1:]
|
|
xy_attn_mask = xy_attn_mask[:, -1:]
|
|
|
|
out, present = self.decoder((xy_input, None), mask=xy_attn_mask, past=past)
|
|
if isinstance(out, tuple): # get rid of stage_embedding
|
|
out = out[0]
|
|
|
|
if out.shape[1] > x_lens.max(): # the first pass, not kvcache yet
|
|
return out[:, x_lens.max():], present
|
|
else: # used kvcache
|
|
return out, present
|
|
|
|
def forward(self, batch):
|
|
"""
|
|
Args:
|
|
x:
|
|
A 2-D tensor of shape (N, S).
|
|
x_lens:
|
|
A 1-D tensor of shape (N,). It contains the number of tokens in `x`
|
|
before padding.
|
|
y:
|
|
A 3-D tensor of shape (N, K, T).
|
|
where K is the number of codebooks
|
|
y_lens:
|
|
A 1-D tensor of shape (N,). It contains the number of tokens in `x`
|
|
before padding.
|
|
"""
|
|
x, x_lens, y, y_lens = batch["x"], batch["x_lens"], batch["y"], batch["y_lens"]
|
|
x = x[:, :x_lens.max()] # this deal with gradient accumulation, where x_lens.max() might not be longer than the length of the current slice of x
|
|
y = y[:, :y_lens.max()]
|
|
assert x.ndim == 2, x.shape
|
|
assert x_lens.ndim == 1, x_lens.shape
|
|
assert y.ndim == 3 and y.shape[1] == self.args.n_codebooks, y.shape
|
|
assert y_lens.ndim == 1, y_lens.shape
|
|
# makes attention mask and padding mask for x
|
|
x_padding_mask = make_pad_mask(x_lens).to(x.device)
|
|
x_attention_mask = torch.triu(torch.ones(x.shape[1], x.shape[1]), diagonal=1).bool().to(x_padding_mask.device)
|
|
x_input = self.text_embedding(x)
|
|
x_input = self.text_positional_embedding(x_input)
|
|
y_input, new_y_lens, targets, y_padding_mask, y_attention_mask, mask_position, patterns = self.prepare_input_target(y, y_lens)
|
|
y_out = self.dec_forward(
|
|
x_input,
|
|
x_lens,
|
|
x_attention_mask,
|
|
x_padding_mask,
|
|
y_input,
|
|
new_y_lens,
|
|
y_attention_mask,
|
|
y_padding_mask
|
|
)
|
|
y_out = y_out[0] # no kv-caching during training
|
|
assert y_out.shape == y_input.shape, f"y_out.shape: {y_out.shape}, y_input.shape: {y_input.shape}" # [B S D]
|
|
|
|
logits = torch.stack([self.predict_layer[i](y_out) for i in range(self.args.n_codebooks)], dim=1) # [B K S card]
|
|
# take out the mask token (using mask_position and new_y_lens) and revert (using function provided by self.pattern)
|
|
assert logits.shape[1] == self.args.n_codebooks and logits.shape[3] == self.n_audio_tokens[0], logits.shape
|
|
|
|
logits_use = self.remove_mask(logits, mask_position, new_y_lens)
|
|
|
|
# revert the pattern shift for each logits section in each sample
|
|
logits_final, logit_masks = self.revert_pattern(patterns, logits_use)
|
|
assert logits_final[0][0].shape[0] == self.args.n_codebooks and logits_final[0][0].shape[2] == self.n_audio_tokens[0], f"it is: {logits_final[0][0].shape}, but should be [K, T, card]"
|
|
# testing
|
|
sample_to_test = 0
|
|
assert len(logits_final[sample_to_test]) == len(targets[sample_to_test]), f"{len(logits_final[sample_to_test])}, {len(targets[sample_to_test])}"
|
|
temp = sum([logits_final[sample_to_test][i].shape[:-1] != targets[sample_to_test][i].shape for i in range(len(targets[sample_to_test]))])
|
|
assert temp == 0, f"none equal positions: {temp}, total number of elements: {len(targets[sample_to_test])}"
|
|
|
|
logit_masked = sum([(item==False).any() for cur_mask in logit_masks for item in cur_mask])
|
|
assert logit_masked == 0, logit_masks
|
|
|
|
logits = torch.cat([torch.cat(item, dim=1) for item in logits_final], dim=1) # [K, T1+T2+T3+..., card]
|
|
targets = torch.cat([torch.cat(item, dim=1) for item in targets], dim=1) # [K, T1+T2+T3+...]
|
|
assert targets.shape[0] == logits.shape[0], f"{targets.shape}, {logits.shape}"
|
|
loss = []
|
|
ntokens = []
|
|
top10acc = []
|
|
for k, (logit, target) in enumerate(zip(logits, targets)):
|
|
loss.append(F.cross_entropy(logit, target, reduction='mean'))
|
|
top10acc.append(self.accuracy_metrics[k](logit.detach(), target))
|
|
ntokens.append(len(logit))
|
|
|
|
all_ntokens = sum(ntokens)
|
|
if self.args.codebook_weight != None:
|
|
codebook_weight = eval(self.args.codebook_weight)
|
|
else:
|
|
codebook_weight = [1.] * self.args.n_codebooks
|
|
loss = sum([l*nt*cw for l, nt, cw in zip(loss, ntokens, codebook_weight)])
|
|
top10acc_by_codebook = [t10a*nt for t10a, nt in zip(top10acc, ntokens)]
|
|
top10acc = sum(top10acc_by_codebook)
|
|
ntokens = torch.tensor(all_ntokens).to(logits.device)
|
|
|
|
return {
|
|
"loss": loss,
|
|
"top10acc": top10acc,
|
|
"top10acc_by_codebook": top10acc_by_codebook,
|
|
"effective_ntoken": ntokens,
|
|
}
|
|
|
|
def inference(
|
|
self,
|
|
x: torch.Tensor,
|
|
x_lens: torch.Tensor,
|
|
y: torch.Tensor,
|
|
mask_interval: list[torch.Tensor],
|
|
top_k: int=-100,
|
|
top_p: float=1.0,
|
|
temperature: float=1.0,
|
|
stop_repetition: int=-1,
|
|
kvcache: int=1,
|
|
silence_tokens: list[int]=[1388,1898,131],
|
|
) -> torch.Tensor:
|
|
"""
|
|
Args:
|
|
x:
|
|
A 2-D tensor of shape (1, L).
|
|
x_lens:
|
|
A 1-D tensor of shape (1,). It contains the number of tokens in `x`
|
|
before padding.
|
|
y:
|
|
A 3-D tensor of shape (1, T, K).
|
|
mask_interval:
|
|
a list of tensors of shape (M, 2). contains M mask_start and mask_end. list length is actually 1, because we only support single sample inference for now
|
|
top_k: (`optional`) int
|
|
The number of highest probability tokens to keep for top-k-filtering. Default to -100.
|
|
top_p: (`optional`) float
|
|
For Neucleus sampling
|
|
temperature: (`optional`) float
|
|
The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
|
|
eog_coef: (`optional`) float
|
|
if 0, no change to eog token logits, otherwise, will adjust eog token logit based on the difference between acoustic token and phn token length
|
|
stop_repetition (`optional`) int
|
|
if not -1, will set the logits of a token that repeated this many times to be -100000, to avoid generating it again. This only apply to tokens from the first codebook
|
|
allowed_repeat_tokens (`optional`) list of ints
|
|
by inspecting the validation set, get a few tokens that indeed repeat a significant amount of time, and exclude those tokens from prevent repetition
|
|
ultimate_stop_repetition (`optional`) int
|
|
no matter that token it is, stop repetition once after this number
|
|
"""
|
|
assert x.ndim == 2, x.shape
|
|
assert x_lens.ndim == 1, x_lens.shape
|
|
assert y.ndim == 3, y.shape
|
|
if self.args.special_first:
|
|
y = y + int(self.args.n_special)
|
|
y = y.transpose(2,1) # [1,T,K] -> [1,K,T]
|
|
assert y.shape[0] == 1 and y.shape[1] == self.args.n_codebooks, y.shape # there is no padding
|
|
assert mask_interval.shape == torch.Size((1, mask_interval.shape[1], 2)), mask_interval
|
|
|
|
# make x attention mask and x_input
|
|
x_attention_mask = torch.triu(torch.ones(x.shape[1], x.shape[1]), diagonal=1).bool().to(x.device)
|
|
# x_attention_mask = torch.zeros(x.shape[1], x.shape[1]).bool().to(x.device)
|
|
x_input = self.text_embedding(x)
|
|
x_input = self.text_positional_embedding(x_input)
|
|
|
|
# make initial y_input
|
|
|
|
# make mask_interval and non_mask_interval
|
|
y_len = y.shape[2]
|
|
y_lens = torch.LongTensor([y_len]).to(y.device)
|
|
mask_interval = mask_interval[0]
|
|
starts = [item[0].item() for item in mask_interval] + [y_len]
|
|
ends = [0] + [item[1].item() for item in mask_interval]
|
|
mask_intervals = [[
|
|
(item[0].item(), item[1].item()) for item in mask_interval
|
|
]] # a werid name change, mask_interval is input, now is mask_intervals, with one more dimension
|
|
non_mask_intervals = [[
|
|
(ns, ne) for ns, ne in zip(ends, starts)
|
|
]]
|
|
|
|
# rearrange y
|
|
# will add have EOG in each section (SOG will be generated by the pattern class)
|
|
# but mask can be inserted later after we have shifted the input
|
|
# y could be rearranged in this way:
|
|
# [
|
|
# [tensor[4, 12], tensor[4, 45], tensor[4, 102], tensor[4, 32]], tensor[4, 22]],
|
|
# [tensor[4, 44], tensor[4, 56], tensor[4, 19]],
|
|
# ...
|
|
# ]
|
|
# for the first list of tensors (4 tensors), first 3 tensors are non_masked part, last 2 are masked part.
|
|
# NOTE #non_masked_part = #masked_part + 1
|
|
rearranged_y = self.rearrange(y, non_mask_intervals, mask_intervals)
|
|
assert rearranged_y[0][0].shape[0] == self.args.n_codebooks, rearranged_y[0][0].shape
|
|
|
|
# shift each element of y
|
|
# next we need to apply pattern shifting to each tensor, after which, we'll replace the starting tokens of each section with a token that's different from the special padding token
|
|
# [
|
|
# [empty, 1, 2, 3, eog, empty, empty, empty],
|
|
# [empty, empty, 1, 2, 3, eog, empty, empty],
|
|
# [empty, empty, empty, 1, 2, 3, eog, empty],
|
|
# [empty, empty, empty, empty, 1, 2, 3, eog]
|
|
# ]
|
|
shifted_y, patterns = self.shift(rearranged_y) # each element [K S], patterns is not used, as we directly use the original input y
|
|
assert shifted_y[0][0].shape[0] == self.args.n_codebooks, shifted_y[0][0].shape
|
|
|
|
# insert mask token at the intersction of each tensor, but *actually inserted eog as place holder*
|
|
# the position of inserted mask is also recorded
|
|
# and the mask_value, the index of the mask emb is recorded
|
|
inserted_y, mask_position, mask_value = self.insert_mask(shifted_y)
|
|
assert inserted_y[0][0].shape[0] == self.args.n_codebooks, inserted_y[0][0].shape[0]
|
|
assert inserted_y[0][1].shape == torch.Size((self.args.n_codebooks, 1)), f"this should be a mask, so should have shape {(self.args.n_codebooks, 1)}, but it's {inserted_y[0][1].shape}"
|
|
|
|
# then concat tensors that belong to the same sample (in order) then get the length of each sample, and then stack them in batch dimension, pad them with pad_token
|
|
cated_y, new_y_lens = self.cat_y(inserted_y, mask_position, y_lens) # KTB
|
|
assert cated_y.shape == torch.Size((self.args.n_codebooks, cated_y.shape[1], len(inserted_y)))
|
|
assert not (cated_y == self.args.audio_pad_token).any(), cated_y
|
|
|
|
### NOTE this is different from forward, as we will remove the masked tokens
|
|
### say there are two masked region
|
|
### the cated_y should be like
|
|
### [empty a a a a mask0 empty b b b mask1 empty c c mask0 empty]
|
|
### which means we need to take the part after the last empty out
|
|
num_mask = len(mask_position[0])//2
|
|
assert num_mask == len(mask_position[0])/2, mask_position
|
|
cated_y = cated_y[:, :mask_position[0][num_mask]+2] # of shape [K,T,B]
|
|
# logging.info(f"mask_position[0][num_mask]+2: {mask_position[0][num_mask]+2}")
|
|
more_mask_value = mask_value[0][num_mask+1:] # NOTE this will be used in the generation loop for reference for inserting mask embedding
|
|
new_y_lens[0] = mask_position[0][num_mask]+2
|
|
mask_position[0] = mask_position[0][:num_mask+1]
|
|
assert mask_position[0][num_mask]+2 == cated_y.shape[1], f"num_mask: {num_mask}, mask_position: {mask_position}, cated_y.shape: {cated_y.shape}"
|
|
|
|
# embed: remember to separately embed the mask tokens
|
|
embedded_y = self.embed_y(cated_y, mask_position, [mask_value[0][:num_mask+1]]) #BTD
|
|
# assert embedded_y.shape == torch.Size((y.shape[0], max(new_y_lens), self.args.d_model)), embedded_y.shape
|
|
|
|
# positional embedding
|
|
y_input = self.audio_positional_embedding(embedded_y)
|
|
|
|
# make attention mask and padding mask
|
|
y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y.device)
|
|
# y_lens = torch.LongTensor([y_input.shape[1]]).to(y.device)
|
|
|
|
x_padding_mask = torch.full((1,x_lens[0]), False).to(x.device)
|
|
y_padding_mask = torch.full((1,new_y_lens[0]), False).to(y.device)
|
|
|
|
|
|
codebook_eog = [False] * self.args.n_codebooks
|
|
generated = [] # doesn't contain any empty_token, contains eog
|
|
cur_generated = []
|
|
# say 0 is empty, 4 is eog
|
|
# tensor([[ 1, 2, 3, 4, 0, 0],
|
|
# [ 0, 1, 2, 3, 4, 0],
|
|
# [ 0, 0, 1, 2, 3, 4]])
|
|
num_gen = []
|
|
cur_num_gen = 0
|
|
##################### silence repetition handling #####################
|
|
##################### silence repetition handling #####################
|
|
logging.info(f"silence tokens: {silence_tokens}, note that if you are not using the pretrained encodec 6f79c6a8, make sure you specified it yourself, rather than using the default")
|
|
consec_silence_count = 0
|
|
prev_token = None
|
|
##################### silence repetition handling #####################
|
|
##################### silence repetition handling #####################
|
|
# prepare the cache placeholder
|
|
# n_layers, 2, bsz, num_heads, src_len, head_dim
|
|
past = torch.ones([self.args.num_decoder_layers, 2, x.shape[0]], device=x.device, dtype=torch.float32) if kvcache else None
|
|
# handle multi-span kv-cache
|
|
new_masked_span = False
|
|
|
|
def sample_helper(n_eog, logits, codebook_eog, top_k, top_p, temperature, prev_token, consec_silence_count, stop_repetition, silence_tokens, cur_num_gen):
|
|
if n_eog == 0:
|
|
logits_adjust = logits
|
|
for jj in range(1,self.args.n_codebooks):
|
|
logits_adjust[jj][self.args.eog] = -10000
|
|
logits_adjust[jj][self.args.empty_token] = -10000
|
|
##################### silence repetition handling #####################
|
|
if stop_repetition > 0 and prev_token in silence_tokens and consec_silence_count > stop_repetition:
|
|
if logits_adjust[0, prev_token] < 0:
|
|
logits_adjust[0, prev_token] = logits_adjust[0, prev_token] * (consec_silence_count - (stop_repetition-1))
|
|
else:
|
|
logits_adjust[0, prev_token] = logits_adjust[0, prev_token] / (consec_silence_count - (stop_repetition-1))
|
|
##################### silence repetition handling #####################
|
|
if type(logits_adjust) == list:
|
|
samples_list= []
|
|
for logit in logits_adjust:
|
|
# print(logit)
|
|
# print(logit.shape)
|
|
cur_sample = topk_sampling(
|
|
logit.unsqueeze(0), top_k=top_k, top_p=top_p, temperature=temperature
|
|
) # [1, 1]
|
|
samples_list.append(cur_sample)
|
|
samples = torch.cat(samples_list, dim=0) # [K, 1]
|
|
else:
|
|
samples = topk_sampling(
|
|
logits_adjust, top_k=top_k, top_p=top_p, temperature=temperature
|
|
) # [K, 1]
|
|
assert samples.shape == torch.Size((self.args.n_codebooks, 1)), f"samples.shape: {samples.shape}"
|
|
if cur_num_gen < self.args.n_codebooks-1:
|
|
for jj in range(1, self.args.n_codebooks - cur_num_gen):
|
|
samples[-jj, 0] = self.args.empty_token
|
|
|
|
if (
|
|
samples[0,0] == self.args.eog or torch.argmax(logits[0], dim=-1) == self.args.eog or y_input.shape[1] > x_lens[0] * 10
|
|
): # last one means y is already too long, shouldn't happen, but put it here
|
|
samples[0,0] = self.args.eog
|
|
codebook_eog[0] = True
|
|
##################### silence repetition handling #####################
|
|
##################### silence repetition handling #####################
|
|
if samples[0,0] in silence_tokens and samples[0,0] == prev_token:
|
|
consec_silence_count += 1
|
|
else:
|
|
consec_silence_count = 0
|
|
prev_token = samples[0,0]
|
|
##################### silence repetition handling #####################
|
|
##################### silence repetition handling #####################
|
|
return samples, codebook_eog, prev_token, consec_silence_count
|
|
else:
|
|
assert sum(codebook_eog[i] for i in range(n_eog)) == n_eog, f"codebook_eog: {codebook_eog}, but n_eog: {n_eog}"
|
|
logits_adjust = logits
|
|
for jj in range(n_eog+1,self.args.n_codebooks):
|
|
logits_adjust[jj][self.args.eog] = -10000
|
|
logits_adjust[jj][self.args.empty_token] = -10000
|
|
if type(logits_adjust) == list:
|
|
samples_list= []
|
|
for logit in logits_adjust:
|
|
cur_sample = topk_sampling(
|
|
logit.unsqueeze(0), top_k=top_k, top_p=top_p, temperature=temperature
|
|
) # [1, 1]
|
|
samples_list.append(cur_sample)
|
|
samples = torch.cat(samples_list, dim=0) # [K, 1]
|
|
else:
|
|
samples = topk_sampling(
|
|
logits_adjust, top_k=top_k, top_p=top_p, temperature=temperature
|
|
) # [K, 1]
|
|
for jj in range(n_eog):
|
|
samples[jj, 0] = self.args.empty_token
|
|
samples[n_eog, 0] = self.args.eog
|
|
codebook_eog[n_eog] = True
|
|
return samples, codebook_eog, prev_token, consec_silence_count
|
|
|
|
while True:
|
|
y_out, present = self.dec_forward(
|
|
x_input,
|
|
x_lens,
|
|
x_attention_mask,
|
|
x_padding_mask,
|
|
y_input,
|
|
new_y_lens,
|
|
y_attention_mask,
|
|
y_padding_mask,
|
|
past=past,
|
|
last_3_tokens = new_masked_span
|
|
)
|
|
if new_masked_span:
|
|
new_masked_span = False
|
|
|
|
if past != None:
|
|
past = torch.cat([past, present.to(past.dtype)], dim=-2) if past.ndim > 3 else present.to(past.dtype)
|
|
|
|
y_out = y_out[:, -1:] # only take the last one
|
|
|
|
logits = torch.stack([self.predict_layer[i](y_out) for i in range(self.args.n_codebooks)], dim=1) # [B K S card], B==S==1, so [1 K 1 card]
|
|
logits = logits.squeeze(0).squeeze(1) # [K card]
|
|
assert logits.shape == torch.Size((self.args.n_codebooks, self.n_audio_tokens[0])), f"{logits.shape}"
|
|
|
|
n_eog = sum(codebook_eog)
|
|
assert n_eog < self.args.n_codebooks
|
|
if self.args.eos > 0: # eos stands for end-of-sentence, which shouldn't be used as we are doing speech editing
|
|
for jj in range(self.args.n_codebooks):
|
|
logits[jj][self.args.eos] = -10000.
|
|
# need to use a helper function to hand different n_eog cases
|
|
samples, codebook_eog, prev_token, consec_silence_count = sample_helper(n_eog, logits, codebook_eog, top_k, top_p, temperature, prev_token, consec_silence_count, stop_repetition, silence_tokens, cur_num_gen)
|
|
cur_num_gen += 1
|
|
cur_generated.append(samples.squeeze(-1)) # [K,1] -> [K]
|
|
# get samples_emb
|
|
samples_emb = torch.stack([self.audio_embedding[k](samples[k]) for k in range(self.args.n_codebooks)], dim=0) # [K,1,D]
|
|
samples_emb = samples_emb.sum(dim=0,keepdim=True) # [1,1,D]
|
|
|
|
if sum(codebook_eog) == self.args.n_codebooks: # generation for the current span is done
|
|
# re-init
|
|
codebook_eog = [False] * self.args.n_codebooks
|
|
num_gen.append(cur_num_gen)
|
|
cur_num_gen = 0
|
|
generated.append(cur_generated)
|
|
cur_generated = []
|
|
|
|
# if the current mask span is the last span, then all done
|
|
# else
|
|
# append the next mask token and the four empty tokens to start the next generation
|
|
if len(more_mask_value) > 0:
|
|
next_mask_ind = more_mask_value.pop(0)
|
|
mask_emb = self.mask_embedding[next_mask_ind].unsqueeze(0).unsqueeze(0) # [1,1,D]
|
|
assert mask_emb.shape == torch.Size((1,1,self.args.d_model)), mask_emb.shape
|
|
empty_token = torch.LongTensor([self.args.empty_token]).to(y.device)
|
|
empty_emb = torch.stack([
|
|
self.audio_embedding[k](empty_token) for k in range(self.args.n_codebooks)], dim=0
|
|
).sum(dim=0, keepdim=True) # [1,1,D]
|
|
assert empty_emb.shape == torch.Size((1,1,self.args.d_model)), empty_emb.shape
|
|
extra_emb = torch.cat([mask_emb, empty_emb], dim=1) # [1,2,D]
|
|
samples_emb = torch.cat([samples_emb, extra_emb], dim=1) # [1,3,D] # prev_last_token, mask_token, empty token
|
|
assert samples_emb.shape == torch.Size((1,3,self.args.d_model)), f"samples_emb.shape: {samples_emb.shape}"
|
|
##################### silence repetition handling #####################
|
|
##################### silence repetition handling #####################
|
|
consec_silence_count = 0
|
|
prev_token = None
|
|
##################### silence repetition handling #####################
|
|
##################### silence repetition handling #####################
|
|
|
|
# handling kv-caching for multi-span editing
|
|
new_masked_span = True
|
|
else:
|
|
break
|
|
else:
|
|
assert samples_emb.shape == torch.Size((1,1,self.args.d_model)), f"samples_emb.shape: {samples_emb.shape}"
|
|
|
|
embedded_y = torch.cat([embedded_y, samples_emb], dim=1)
|
|
# positional embedding
|
|
y_input = self.audio_positional_embedding(embedded_y) # [B T D]
|
|
# make attention mask and padding mask
|
|
y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y.device)
|
|
new_y_lens = torch.LongTensor([y_input.shape[1]]).to(y.device)
|
|
y_padding_mask = torch.full((1,new_y_lens[0]), False).to(y.device)
|
|
|
|
assert len(generated) == num_mask, f"len(generated): {len(generated)}, num_mask: {num_mask}"
|
|
|
|
# # combine non_masked_span with generated spans
|
|
# first need to shift the generated part back
|
|
flatten_gen = []
|
|
for l, orig_span in enumerate(generated):
|
|
span = torch.stack(orig_span, dim=0) # [T K]
|
|
span = span.transpose(1,0) # [K, T]
|
|
assert span.shape[0] == self.args.n_codebooks, span.shape
|
|
unshifted_span = []
|
|
for j, s in enumerate(span):
|
|
start_from = j
|
|
end_at = - (self.args.n_codebooks - start_from)
|
|
unshifted_span.append(s[start_from:end_at])
|
|
unshifted_span = torch.stack(unshifted_span, dim=0)
|
|
|
|
assert unshifted_span.shape[1] == num_gen[l] - self.args.n_codebooks, f"len(unshifted_spans[0]): {len(unshifted_span[0])}, num_gen[l]: {num_gen[l]}"
|
|
flatten_gen.append(unshifted_span)
|
|
# logging.info(f"unshfited_span: {unshifted_span.shape}")
|
|
# raise
|
|
assert len(non_mask_intervals[0]) - 1 == len(flatten_gen), f"len(non_mask_intervals[0]): {len(non_mask_intervals[0])}, len(flatten_gen): {len(flatten_gen)}"
|
|
res = []
|
|
for orig_interval, gen in zip(non_mask_intervals[0], flatten_gen):
|
|
res.append(y[0, :, orig_interval[0]:orig_interval[1]])
|
|
res.append(gen)
|
|
res.append(y[0, :, non_mask_intervals[0][-1][0]:non_mask_intervals[0][-1][1]])
|
|
res = torch.cat(res, dim=1).unsqueeze(0) # [K,new_T] -> [1, K, new_T]
|
|
|
|
expected_y_len = y_len - sum([item[1] - item[0] for item in mask_intervals[0]]) + sum([item - self.args.n_codebooks for item in num_gen])
|
|
assert res.shape == torch.Size((1, self.args.n_codebooks, expected_y_len)), f"res.shape: {res.shape}, expected_y_len: {expected_y_len}. y_len - sum([item[1] - item[0] for item in mask_interval]) + sum([item - self.args.n_codebooks for item in num_gen]): {y_len}-{sum([item[1] - item[0] for item in mask_interval])} + {sum([item - self.args.n_codebooks for item in num_gen])}"
|
|
|
|
if self.args.special_first:
|
|
res = res - int(self.args.n_special)
|
|
|
|
return res
|
|
|
|
def inference_tts(
|
|
self,
|
|
x: torch.Tensor,
|
|
x_lens: torch.Tensor,
|
|
y: torch.Tensor,
|
|
top_k: int=-100,
|
|
top_p: float=1.0,
|
|
temperature: float=1.0,
|
|
stop_repetition: int=3,
|
|
kvcache: int=1,
|
|
silence_tokens: list[int]=[1388,1898,131],
|
|
*kargs
|
|
) -> torch.Tensor:
|
|
"""
|
|
different from inference_tts, this implementation uses kvcache, which should have significant speed up
|
|
Args:
|
|
x:
|
|
A 2-D tensor of shape (1, L).
|
|
x_lens:
|
|
A 1-D tensor of shape (1,). It contains the number of tokens in `x`
|
|
before padding.
|
|
y:
|
|
A 3-D tensor of shape (1, T, K).
|
|
top_k: (`optional`) int
|
|
The number of highest probability tokens to keep for top-k-filtering. Default to -100.
|
|
top_p: (`optional`) float
|
|
For Neucleus sampling
|
|
temperature: (`optional`) float
|
|
The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
|
|
"""
|
|
eog_inference = self.args.eos if self.args.eos>0 else self.args.eog
|
|
assert x.ndim == 2, x.shape
|
|
assert x_lens.ndim == 1, x_lens.shape
|
|
assert y.ndim == 3, y.shape
|
|
if self.args.special_first:
|
|
y = y + int(self.args.n_special)
|
|
y = y.transpose(2,1) # [1,T,K] -> [1,K,T]
|
|
assert y.shape[0] == 1 and y.shape[1] == self.args.n_codebooks, y.shape # there is no padding
|
|
|
|
# make x attention mask and x_input
|
|
x_attention_mask = torch.triu(torch.ones(x.shape[1], x.shape[1]), diagonal=1).bool().to(x.device)
|
|
# x_attention_mask = torch.zeros(x.shape[1], x.shape[1]).bool().to(x.device)
|
|
x_input = self.text_embedding(x)
|
|
x_input = self.text_positional_embedding(x_input)
|
|
|
|
y_len = y.shape[2]
|
|
y_lens = torch.LongTensor([y_len]).to(y.device)
|
|
|
|
# rearrange y, we don't add eog to the end, this doesn't actually do anything in the tts scenario
|
|
rearranged_y = [[y[0]]]
|
|
assert rearranged_y[0][0].shape[0] == self.args.n_codebooks, rearranged_y[0][0].shape
|
|
|
|
# shift y to create the delayed pattern
|
|
shifted_y, patterns = self.shift(rearranged_y) # each element [K S], patterns is not used, as we directly use the original input y
|
|
assert shifted_y[0][0].shape[0] == self.args.n_codebooks, shifted_y[0][0].shape
|
|
assert len(shifted_y[0]) == 1, len(shifted_y[0])
|
|
|
|
# below is different from forward or inference
|
|
# where we cut this shifted part
|
|
shifted_y[0][0] = shifted_y[0][0][:, :-(self.args.n_codebooks-1)]
|
|
assert not (shifted_y[0][0][self.args.n_codebooks:] == self.args.empty_token).any() and not (shifted_y[0][0][self.args.n_codebooks:] == self.args.eog).any(), shifted_y[0][0]
|
|
|
|
# next section in inference is insert mask at the intersection of each tensor in a sample, but we don't need to do that
|
|
# next section is concate tensors of each sample to one tensor, which we also don't need
|
|
cated_y = shifted_y[0][0].unsqueeze(-1) #[K,S]->[K,S,B]
|
|
new_y_lens = torch.LongTensor([cated_y.shape[1]]).to(cated_y.device)
|
|
assert cated_y.shape == torch.Size((self.args.n_codebooks, cated_y.shape[1], 1))
|
|
assert not (cated_y == self.args.audio_pad_token).any(), cated_y
|
|
|
|
# replace tokens in y with the embeddings, add sum codebooks up
|
|
embedded_y = torch.stack([self.audio_embedding[k](cated_y[k]) for k in range(self.args.n_codebooks)], dim=0) # [K, S, B, D]
|
|
assert embedded_y.shape[0] == self.args.n_codebooks, embedded_y.shape
|
|
assert embedded_y.shape[-1] == self.args.d_model, embedded_y.shape
|
|
embedded_y = embedded_y.sum(dim=0) # [K,S,B,D]->[S,B,D]
|
|
embedded_y = embedded_y.transpose(1,0) # [S,B,D]->[B,S,D]
|
|
|
|
# positional embedding
|
|
y_input = self.audio_positional_embedding(embedded_y)
|
|
|
|
# make attention mask and padding mask
|
|
y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y.device)
|
|
|
|
x_padding_mask = torch.full((1,x_lens[0]), False).to(x.device)
|
|
y_padding_mask = torch.full((1,new_y_lens[0]), False).to(y.device)
|
|
|
|
# entering the generation stage
|
|
# starting from line 708
|
|
codebook_eog = [False] * self.args.n_codebooks
|
|
generated = [] # doesn't contain any empty token, contain eog
|
|
cur_generated = []
|
|
# say 0 is empty, 4 is eog
|
|
# tensor([[ 1, 2, 3, 4, 0, 0],
|
|
# [ 0, 1, 2, 3, 4, 0],
|
|
# [ 0, 0, 1, 2, 3, 4]])
|
|
num_gen = []
|
|
cur_num_gen = 0
|
|
##################### silence repetition handling #####################
|
|
##################### silence repetition handling #####################
|
|
logging.info(f"silence tokens: {silence_tokens}, note that if you are not using the pretrained encodec 6f79c6a8, make sure you specified it yourself, rather than using the default")
|
|
consec_silence_count = 0
|
|
prev_token = None
|
|
##################### silence repetition handling #####################
|
|
##################### silence repetition handling #####################
|
|
|
|
# prepare the cache placeholder
|
|
# n_layers, 2, bsz, num_heads, src_len, head_dim
|
|
past = torch.ones([self.args.num_decoder_layers, 2, x.shape[0]], device=x.device, dtype=torch.float32) if kvcache else None
|
|
# logging.info(f"number of decoder layers: {self.args.num_decoder_layers}")
|
|
# logging.info(f"number of decoder layers: {self.args.num_decoder_layers}")
|
|
# logging.info(f"number of decoder layers: {self.args.num_decoder_layers}")
|
|
def sample_helper(n_eog, logits, codebook_eog, top_k, top_p, temperature, prev_token, consec_silence_count, stop_repetition, silence_tokens, cur_num_gen):
|
|
if n_eog == 0:
|
|
logits_adjust = logits
|
|
for jj in range(1,self.args.n_codebooks):
|
|
logits_adjust[jj][eog_inference] = -10000
|
|
logits_adjust[jj][self.args.empty_token] = -10000
|
|
if cur_num_gen <= self.args.encodec_sr // 5: # this shouldn't happen, but just in case the model stopped too early
|
|
logits_adjust[0][eog_inference] = -10000
|
|
##################### silence repetition handling #####################
|
|
if stop_repetition > 0 and prev_token in silence_tokens and consec_silence_count > stop_repetition:
|
|
if logits_adjust[0, prev_token] < 0:
|
|
logits_adjust[0, prev_token] = logits_adjust[0, prev_token] * (consec_silence_count - (stop_repetition-1))
|
|
else:
|
|
logits_adjust[0, prev_token] = logits_adjust[0, prev_token] / (consec_silence_count - (stop_repetition-1))
|
|
##################### silence repetition handling #####################
|
|
samples = topk_sampling(
|
|
logits_adjust, top_k=top_k, top_p=top_p, temperature=temperature
|
|
) # [K, 1]
|
|
assert samples.shape == torch.Size((self.args.n_codebooks, 1)), f"samples.shape: {samples.shape}"
|
|
if cur_num_gen < self.args.n_codebooks-1:
|
|
for jj in range(1, self.args.n_codebooks - cur_num_gen):
|
|
samples[-jj, 0] = self.args.empty_token
|
|
|
|
if (
|
|
samples[0,0] == eog_inference or torch.argmax(logits[0], dim=-1) == eog_inference or y_input.shape[1] > x_lens[0] * (self.args.encodec_sr//5)
|
|
): # last one means y is already too long, shouldn't happen, but put it here
|
|
samples[0,0] = eog_inference
|
|
codebook_eog[0] = True
|
|
##################### silence repetition handling #####################
|
|
if samples[0,0] in silence_tokens and samples[0,0] == prev_token:
|
|
consec_silence_count += 1
|
|
else:
|
|
consec_silence_count = 0
|
|
prev_token = samples[0,0]
|
|
##################### silence repetition handling #####################
|
|
return samples, codebook_eog, prev_token, consec_silence_count
|
|
else:
|
|
assert sum(codebook_eog[i] for i in range(n_eog)) == n_eog, f"codebook_eog: {codebook_eog}, but n_eog: {n_eog}"
|
|
logits_adjust = logits
|
|
for jj in range(n_eog+1,self.args.n_codebooks):
|
|
logits_adjust[jj][eog_inference] = -10000
|
|
logits_adjust[jj][self.args.empty_token] = -10000
|
|
samples = topk_sampling(
|
|
logits_adjust, top_k=top_k, top_p=top_p, temperature=temperature
|
|
) # [K, 1]
|
|
for jj in range(n_eog):
|
|
samples[jj, 0] = self.args.empty_token
|
|
samples[n_eog, 0] = eog_inference
|
|
codebook_eog[n_eog] = True
|
|
return samples, codebook_eog, prev_token, consec_silence_count
|
|
while True:
|
|
y_out, present = self.dec_forward(
|
|
x_input,
|
|
x_lens,
|
|
x_attention_mask,
|
|
x_padding_mask,
|
|
y_input,
|
|
new_y_lens,
|
|
y_attention_mask,
|
|
y_padding_mask,
|
|
past=past
|
|
)
|
|
if past != None:
|
|
past = torch.cat([past, present.to(past.dtype)], dim=-2) if past.ndim > 3 else present.to(past.dtype)
|
|
|
|
|
|
y_out = y_out[:, -1:] # only take the last token
|
|
logits = torch.stack([self.predict_layer[i](y_out) for i in range(self.args.n_codebooks)], dim=1) # [B K S card], B==S==1, so [1 K 1 card]
|
|
logits = logits.squeeze(0).squeeze(1) # [K card]
|
|
assert logits.shape == torch.Size((self.args.n_codebooks, self.n_audio_tokens[0])), f"{logits.shape}"
|
|
|
|
n_eog = sum(codebook_eog)
|
|
assert n_eog < self.args.n_codebooks
|
|
if self.args.eos > 0: # if we are using end-of-sentence token (which is used by default), eog shouldn't be used here, as there is no masked spans
|
|
for jj in range(self.args.n_codebooks):
|
|
logits[jj][self.args.eog] = -10000.
|
|
|
|
samples, codebook_eog, prev_token, consec_silence_count = sample_helper(n_eog, logits, codebook_eog, top_k, top_p, temperature, prev_token, consec_silence_count, stop_repetition, silence_tokens, cur_num_gen)
|
|
|
|
cur_num_gen += 1
|
|
cur_generated.append(samples.squeeze(-1)) # [K,1] -> [K]
|
|
|
|
# samples.shape is [K,1]
|
|
# ge samples_emb
|
|
samples_emb = torch.stack([self.audio_embedding[k](samples[k]) for k in range(self.args.n_codebooks)], dim=0) # [K,1,D]
|
|
samples_emb = samples_emb.sum(dim=0,keepdim=True) # [1,1,D]
|
|
|
|
if sum(codebook_eog) == self.args.n_codebooks: # generation for the current span is done
|
|
codebook_eog = [False] * self.args.n_codebooks
|
|
num_gen.append(cur_num_gen)
|
|
cur_num_gen = 0
|
|
generated.append(cur_generated)
|
|
cur_generated = []
|
|
break
|
|
else:
|
|
assert samples_emb.shape == torch.Size((1,1,self.args.d_model)), f"samples_emb.shape: {samples_emb.shape}"
|
|
|
|
embedded_y = torch.cat([embedded_y, samples_emb], dim=1)
|
|
y_input = self.audio_positional_embedding(embedded_y) # [B T D]
|
|
# make attention mask and padding mask
|
|
y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y.device)
|
|
new_y_lens = torch.LongTensor([y_input.shape[1]]).to(y.device)
|
|
y_padding_mask = torch.full((1,new_y_lens[0]), False).to(y.device)
|
|
|
|
assert len(generated) == 1, f"len(generated): {len(generated)}"
|
|
|
|
# revert the pattern
|
|
flatten_gen = []
|
|
for l, orig_span in enumerate(generated):
|
|
span = torch.stack(orig_span, dim=0) # [T, K]
|
|
span = span.transpose(1,0) # [K, T]
|
|
assert span.shape[0] == self.args.n_codebooks, span.shape
|
|
unshifted_span = []
|
|
for j, s in enumerate(span):
|
|
start_from = j
|
|
end_at = - (self.args.n_codebooks - start_from)
|
|
unshifted_span.append(s[start_from:end_at])
|
|
unshifted_span = torch.stack(unshifted_span, dim=0)
|
|
|
|
assert unshifted_span.shape[1] == num_gen[l] - self.args.n_codebooks, f"len(unshifted_spans[0]): {len(unshifted_span[0])}, num_gen[l]: {num_gen[l]}"
|
|
|
|
flatten_gen.append(unshifted_span)
|
|
assert len(flatten_gen) == 1, len(flatten_gen)
|
|
|
|
# combine
|
|
res = [y[0], flatten_gen[0]]
|
|
res = torch.cat(res, dim=1).unsqueeze(0) # [K, new_t] -> [1, K, new_T]
|
|
|
|
expected_y_len = y_len + sum([item - self.args.n_codebooks for item in num_gen])
|
|
assert res.shape == torch.Size((1, self.args.n_codebooks, expected_y_len)), f"res.shape: {res.shape}, expected_y_len: {expected_y_len}. y_len + sum([item - self.args.n_codebooks for item in num_gen]): {y_len} + {sum([item - self.args.n_codebooks for item in num_gen])}"
|
|
|
|
if self.args.special_first:
|
|
res = res - int(self.args.n_special)
|
|
flatten_gen = flatten_gen - int(self.args.n_special)
|
|
|
|
return res, flatten_gen[0].unsqueeze(0)
|
|
|
|
|
|
def inference_tts_batch(
|
|
self,
|
|
x: torch.Tensor,
|
|
x_lens: torch.Tensor,
|
|
y: torch.Tensor,
|
|
top_k: int=-100,
|
|
top_p: float=1.0,
|
|
temperature: float=1.0,
|
|
stop_repetition: int=3,
|
|
kvcache: int=1,
|
|
batch_size: int=5,
|
|
silence_tokens: list[int]=[1388,1898,131],
|
|
*kargs
|
|
) -> torch.Tensor:
|
|
"""
|
|
have a batch size when forward passing, but they are equivalant to same example but different random seed, therefore as long as one example generated eog, we can drop all other samlpes
|
|
different from inference_tts, this implementation uses kvcache, which should have significant speed up
|
|
Args:
|
|
x:
|
|
A 2-D tensor of shape (1, L).
|
|
x_lens:
|
|
A 1-D tensor of shape (1,). It contains the number of tokens in `x`
|
|
before padding.
|
|
y:
|
|
A 3-D tensor of shape (1, T, K).
|
|
top_k: (`optional`) int
|
|
The number of highest probability tokens to keep for top-k-filtering. Default to -100.
|
|
top_p: (`optional`) float
|
|
For Neucleus sampling
|
|
temperature: (`optional`) float
|
|
The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
|
|
"""
|
|
eog_inference = self.args.eos if self.args.eos>0 else self.args.eog
|
|
assert x.ndim == 2, x.shape
|
|
assert x_lens.ndim == 1, x_lens.shape
|
|
assert y.ndim == 3, y.shape
|
|
if self.args.special_first:
|
|
y = y + int(self.args.n_special)
|
|
y = y.transpose(2,1) # [1,T,K] -> [1,K,T]
|
|
assert y.shape[0] == 1 and y.shape[1] == self.args.n_codebooks, y.shape # there is no padding
|
|
|
|
# make x attention mask and x_input
|
|
x_attention_mask = torch.triu(torch.ones(x.shape[1], x.shape[1]), diagonal=1).bool().to(x.device)
|
|
# x_attention_mask = torch.zeros(x.shape[1], x.shape[1]).bool().to(x.device)
|
|
x_input = self.text_embedding(x)
|
|
x_input = self.text_positional_embedding(x_input)
|
|
|
|
y_len = y.shape[2]
|
|
y_lens = torch.LongTensor([y_len]).to(y.device)
|
|
|
|
# rearrange y, we don't add eog to the end, this doesn't actually do anything in the tts scenario
|
|
rearranged_y = [[y[0]]]
|
|
assert rearranged_y[0][0].shape[0] == self.args.n_codebooks, rearranged_y[0][0].shape
|
|
|
|
# shift y to create the delayed pattern
|
|
shifted_y, patterns = self.shift(rearranged_y) # each element [K S], patterns is not used, as we directly use the original input y
|
|
assert shifted_y[0][0].shape[0] == self.args.n_codebooks, shifted_y[0][0].shape
|
|
assert len(shifted_y[0]) == 1, len(shifted_y[0])
|
|
|
|
# below is different from forward or inference
|
|
# where we cut this shifted part
|
|
shifted_y[0][0] = shifted_y[0][0][:, :-(self.args.n_codebooks-1)]
|
|
assert not (shifted_y[0][0][self.args.n_codebooks:] == self.args.empty_token).any() and not (shifted_y[0][0][self.args.n_codebooks:] == self.args.eog).any(), shifted_y[0][0]
|
|
|
|
# next section in inference is insert mask at the intersection of each tensor in a sample, but we don't need to do that
|
|
# next section is concate tensors of each sample to one tensor, which we also don't need
|
|
cated_y = shifted_y[0][0].unsqueeze(-1) #[K,S]->[K,S,B]
|
|
new_y_lens = torch.LongTensor([cated_y.shape[1]]).to(cated_y.device)
|
|
assert cated_y.shape == torch.Size((self.args.n_codebooks, cated_y.shape[1], 1))
|
|
assert not (cated_y == self.args.audio_pad_token).any(), cated_y
|
|
|
|
# replace tokens in y with the embeddings, add sum codebooks up
|
|
embedded_y = torch.stack([self.audio_embedding[k](cated_y[k]) for k in range(self.args.n_codebooks)], dim=0) # [K, S, B, D]
|
|
assert embedded_y.shape[0] == self.args.n_codebooks, embedded_y.shape
|
|
assert embedded_y.shape[-1] == self.args.d_model, embedded_y.shape
|
|
embedded_y = embedded_y.sum(dim=0) # [K,S,B,D]->[S,B,D]
|
|
embedded_y = embedded_y.transpose(1,0) # [S,B,D]->[B,S,D]
|
|
|
|
# positional embedding
|
|
y_input = self.audio_positional_embedding(embedded_y)
|
|
|
|
# make attention mask and padding mask
|
|
y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y.device)
|
|
|
|
x_padding_mask = torch.full((1,x_lens[0]), False).to(x.device)
|
|
y_padding_mask = torch.full((1,new_y_lens[0]), False).to(y.device)
|
|
|
|
# entering the generation stage
|
|
# starting from line 708
|
|
codebook_eog = [False] * self.args.n_codebooks
|
|
generated = [] # doesn't contain any empty token, contain eog
|
|
cur_generated = [[] for _ in range(batch_size)]
|
|
# say 0 is empty, 4 is eog
|
|
# tensor([[ 1, 2, 3, 4, 0, 0],
|
|
# [ 0, 1, 2, 3, 4, 0],
|
|
# [ 0, 0, 1, 2, 3, 4]])
|
|
num_gen = []
|
|
cur_num_gen = 0
|
|
##################### silence repetition handling #####################
|
|
##################### silence repetition handling #####################
|
|
logging.info(f"silence tokens: {silence_tokens}, note that if you are not using the pretrained encodec 6f79c6a8, make sure you specified it yourself, rather than using the default")
|
|
consec_silence_counts = [0 for _ in range(batch_size)]
|
|
prev_tokens = [None for _ in range(batch_size)]
|
|
##################### silence repetition handling #####################
|
|
##################### silence repetition handling #####################
|
|
|
|
# prepare the cache placeholder
|
|
# n_layers, 2, bsz, num_heads, src_len, head_dim
|
|
past = torch.ones([self.args.num_decoder_layers, 2, x.shape[0]], device=x.device, dtype=torch.float32) if kvcache else None
|
|
# logging.info(f"number of decoder layers: {self.args.num_decoder_layers}")
|
|
# logging.info(f"number of decoder layers: {self.args.num_decoder_layers}")
|
|
# logging.info(f"number of decoder layers: {self.args.num_decoder_layers}")
|
|
keep = None # NOTE: this very important, tells which sample to keep
|
|
def sample_helper(n_eog, logits, codebook_eog, top_k, top_p, temperature, prev_tokens, consec_silence_counts, stop_repetition, silence_tokens, cur_num_gen, keep):
|
|
if n_eog == 0:
|
|
logits_adjust = logits
|
|
for jj in range(1,self.args.n_codebooks):
|
|
logits_adjust[:,jj,eog_inference] = -10000
|
|
logits_adjust[:,jj,self.args.empty_token] = -10000
|
|
if cur_num_gen <= self.args.encodec_sr // 5: # this shouldn't happen, but just in case the model stopped too early
|
|
logits_adjust[:,:,eog_inference] = -10000
|
|
##################### silence repetition handling #####################
|
|
for b in range(batch_size):
|
|
prev_token = prev_tokens[b]
|
|
consec_silence_count = consec_silence_counts[b]
|
|
if stop_repetition > 0 and prev_token in silence_tokens and consec_silence_count > stop_repetition:
|
|
if logits_adjust[b, 0, prev_token] < 0:
|
|
logits_adjust[b, 0, prev_token] = logits_adjust[b, 0, prev_token] * (consec_silence_count - (stop_repetition-1))
|
|
else:
|
|
logits_adjust[b, 0, prev_token] = logits_adjust[b, 0, prev_token] / (consec_silence_count - (stop_repetition-1))
|
|
##################### silence repetition handling #####################
|
|
samples = topk_sampling(
|
|
logits_adjust.reshape(batch_size * self.args.n_codebooks, logits_adjust.shape[-1]), top_k=top_k, top_p=top_p, temperature=temperature
|
|
) # [B*K, 1]
|
|
samples = samples.reshape(batch_size, self.args.n_codebooks, 1)
|
|
assert samples.shape == torch.Size((batch_size, self.args.n_codebooks, 1)), f"samples.shape: {samples.shape}"
|
|
for b in range(batch_size):
|
|
if cur_num_gen < self.args.n_codebooks-1:
|
|
for jj in range(1, self.args.n_codebooks - cur_num_gen):
|
|
samples[b, -jj, 0] = self.args.empty_token
|
|
|
|
if (
|
|
samples[b,0,0] == eog_inference or torch.argmax(logits[b,0], dim=-1) == eog_inference or y_input.shape[1] > x_lens[b] * (self.args.encodec_sr//5)
|
|
): # last one means y is already too long, shouldn't happen, but put it here
|
|
samples[b,0,0] = eog_inference
|
|
codebook_eog[0] = True
|
|
keep = b # NOTE keep is a very important variable, we only return this one, note that if eog shows up in two samples, keep will be overwritten by the later one (or the last one)
|
|
##################### silence repetition handling #####################
|
|
if samples[b,0,0] in silence_tokens and samples[b,0,0] == prev_tokens[b]:
|
|
consec_silence_counts[b] += 1
|
|
else:
|
|
consec_silence_counts[b] = 0
|
|
prev_tokens[b] = samples[b,0,0]
|
|
##################### silence repetition handling #####################
|
|
return samples, codebook_eog, prev_tokens, consec_silence_counts, keep
|
|
else:
|
|
assert sum(codebook_eog[i] for i in range(n_eog)) == n_eog, f"codebook_eog: {codebook_eog}, but n_eog: {n_eog}"
|
|
logits_adjust = logits
|
|
for jj in range(n_eog+1,self.args.n_codebooks):
|
|
logits_adjust[:,jj,eog_inference] = -10000
|
|
logits_adjust[:,jj,self.args.empty_token] = -10000
|
|
samples = topk_sampling(
|
|
logits_adjust.reshape(batch_size * self.args.n_codebooks, logits_adjust.shape[-1]), top_k=top_k, top_p=top_p, temperature=temperature
|
|
) # [B, K, 1]
|
|
samples = samples.reshape(batch_size, self.args.n_codebooks, 1)
|
|
for jj in range(n_eog):
|
|
samples[keep, jj, 0] = self.args.empty_token
|
|
samples[keep, n_eog, 0] = eog_inference
|
|
codebook_eog[n_eog] = True
|
|
return samples, codebook_eog, prev_tokens, consec_silence_counts, keep
|
|
while True:
|
|
# if cur_num_gen > 0, should have everything in kvcache, so only pass in the last token
|
|
# in the first generation step, we repeat each tensor to make their first dimension of length the batch size
|
|
if cur_num_gen == 0:
|
|
assert x_input.ndim == 3 and x_input.shape[0] == 1, x_input.shape
|
|
assert x_padding_mask.ndim == 2 and x_padding_mask.shape[0] == 1, x_padding_mask.shape
|
|
assert y_input.ndim == 3 and y_input.shape[0] == 1 and y_input.shape[1] == new_y_lens[0], y_input.shape
|
|
assert embedded_y.ndim == 3 and embedded_y.shape[0] == 1 and embedded_y.shape[1] == new_y_lens[0], embedded_y.shape
|
|
x_input = x_input.repeat(batch_size, 1, 1)
|
|
x_lens = x_lens.repeat(batch_size)
|
|
# x_attention_mask = x_attention_mask.repeat(batch_size, 1, 1) # no need to work with attention mask, it doesn't contain batch dimension
|
|
x_padding_mask = x_padding_mask.repeat(batch_size, 1)
|
|
y_input = y_input.repeat(batch_size, 1, 1)
|
|
new_y_lens = new_y_lens.repeat(batch_size)
|
|
# y_attention_mask = y_attention_mask.repeat(batch_size, 1, 1) # no need to work with attention mask, it doesn't contain batch dimension
|
|
y_padding_mask = y_padding_mask.repeat(batch_size, 1)
|
|
embedded_y = embedded_y.repeat(batch_size, 1, 1) # will be used to concat with newly generated token embedding
|
|
past = past.repeat(1, 1, batch_size) if past != None else None
|
|
else:
|
|
assert x_input.shape[0] == batch_size and x_padding_mask.shape[0] == batch_size and y_input.shape[0] == batch_size and new_y_lens.shape[0] == batch_size, f"x_input.shape: {x_input.shape}, x_padding_mask.shape: {x_padding_mask.shape}, y_input.shape: {y_input.shape}, new_y_lens.shape: {new_y_lens.shape}"
|
|
y_out, present = self.dec_forward(
|
|
x_input,
|
|
x_lens,
|
|
x_attention_mask,
|
|
x_padding_mask,
|
|
y_input,
|
|
new_y_lens,
|
|
y_attention_mask,
|
|
y_padding_mask,
|
|
past=past
|
|
)
|
|
if past != None:
|
|
past = torch.cat([past, present.to(past.dtype)], dim=-2) if past.ndim > 3 else present.to(past.dtype)
|
|
|
|
# if no eog emerges, y_out should have batch size of batch_size
|
|
if sum(codebook_eog) == 0:
|
|
assert y_out.shape[0] == batch_size and y_out.ndim == 3, y_out.shape
|
|
y_out = y_out[:, -1:] # only take the last token
|
|
logits = torch.stack([self.predict_layer[i](y_out) for i in range(self.args.n_codebooks)], dim=1) # [B K S card], S==1, so [B K 1 card]
|
|
logits = logits.squeeze(2) # [B K card]
|
|
assert logits.shape == torch.Size((batch_size, self.args.n_codebooks, self.n_audio_tokens[0])), f"{logits.shape}"
|
|
|
|
n_eog = sum(codebook_eog)
|
|
if self.args.eos > 0:
|
|
for jj in range(self.args.n_codebooks):
|
|
logits[:,jj,self.args.eog] = -10000.
|
|
samples, codebook_eog, prev_tokens, consec_silence_counts, keep = sample_helper(n_eog, logits, codebook_eog, top_k, top_p, temperature, prev_tokens, consec_silence_counts, stop_repetition, silence_tokens, cur_num_gen, keep)
|
|
|
|
cur_num_gen += 1
|
|
if sum(codebook_eog) == 0: # no eog yet, keep batch_size of samples
|
|
assert keep == None
|
|
for b in range(batch_size):
|
|
cur_generated[b].append(samples[b].squeeze(-1))
|
|
elif sum(codebook_eog) == 1: # the first eog just showed up in this step
|
|
assert keep != None
|
|
cur_generated = cur_generated[keep]
|
|
cur_generated.append(samples[keep].squeeze(-1))
|
|
else: # we are generating the rest eogs for the 'keep' sample
|
|
cur_generated.append(samples[keep].squeeze(-1))
|
|
|
|
# samples.shape is [K,1]
|
|
# ge samples_emb
|
|
samples_emb = torch.stack([self.audio_embedding[k](samples[:, k]) for k in range(self.args.n_codebooks)], dim=1) # [B, K,1,D]
|
|
assert samples_emb.shape == torch.Size([batch_size, self.args.n_codebooks, 1, self.args.d_model])
|
|
samples_emb = samples_emb.sum(dim=1,keepdim=False) # [B,1,D]
|
|
if sum(codebook_eog) == self.args.n_codebooks: # generation for the current span is done
|
|
codebook_eog = [False] * self.args.n_codebooks
|
|
num_gen.append(cur_num_gen)
|
|
cur_num_gen = 0
|
|
generated.append(cur_generated)
|
|
cur_generated = [[] for _ in range(batch_size)]
|
|
break
|
|
else:
|
|
assert samples_emb.shape == torch.Size((batch_size,1,self.args.d_model)), f"samples_emb.shape: {samples_emb.shape}"
|
|
|
|
embedded_y = torch.cat([embedded_y, samples_emb], dim=1)
|
|
y_input = self.audio_positional_embedding(embedded_y) # [B T D]
|
|
# make attention mask and padding mask
|
|
y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y.device)
|
|
new_y_lens = torch.LongTensor([y_input.shape[1]]).to(y.device).repeat(batch_size)
|
|
y_padding_mask = torch.full((batch_size,new_y_lens[0]), False).to(y.device)
|
|
|
|
assert len(generated) == 1, f"len(generated): {len(generated)}"
|
|
|
|
# revert the pattern
|
|
flatten_gen = []
|
|
for l, orig_span in enumerate(generated):
|
|
span = torch.stack(orig_span, dim=0) # [T, K]
|
|
span = span.transpose(1,0) # [K, T]
|
|
assert span.shape[0] == self.args.n_codebooks, span.shape
|
|
unshifted_span = []
|
|
for j, s in enumerate(span):
|
|
start_from = j
|
|
end_at = - (self.args.n_codebooks - start_from)
|
|
unshifted_span.append(s[start_from:end_at])
|
|
unshifted_span = torch.stack(unshifted_span, dim=0)
|
|
|
|
assert unshifted_span.shape[1] == num_gen[l] - self.args.n_codebooks, f"len(unshifted_spans[0]): {len(unshifted_span[0])}, num_gen[l]: {num_gen[l]}"
|
|
|
|
flatten_gen.append(unshifted_span)
|
|
assert len(flatten_gen) == 1, len(flatten_gen)
|
|
|
|
# combine
|
|
res = [y[0], flatten_gen[0]]
|
|
res = torch.cat(res, dim=1).unsqueeze(0) # [K, new_t] -> [1, K, new_T]
|
|
|
|
expected_y_len = y_len + sum([item - self.args.n_codebooks for item in num_gen])
|
|
assert res.shape == torch.Size((1, self.args.n_codebooks, expected_y_len)), f"res.shape: {res.shape}, expected_y_len: {expected_y_len}. y_len + sum([item - self.args.n_codebooks for item in num_gen]): {y_len} + {sum([item - self.args.n_codebooks for item in num_gen])}"
|
|
|
|
if self.args.special_first:
|
|
res = res - int(self.args.n_special)
|
|
flatten_gen = flatten_gen - int(self.args.n_special)
|
|
|
|
return res, flatten_gen[0].unsqueeze(0) |