more code refactor
This commit is contained in:
Blaise
@@ -1,5 +1,6 @@
|
||||
# modified from https://github.com/feng-yufei/shared_debugging_code/blob/main/model/t2s_lightning_module.py
|
||||
import os,sys
|
||||
import os, sys
|
||||
|
||||
now_dir = os.getcwd()
|
||||
sys.path.append(now_dir)
|
||||
from typing import Dict
|
||||
@@ -12,29 +13,35 @@ from AR.modules.optim import ScaledAdam
|
||||
|
||||
|
||||
class Text2SemanticLightningModule(LightningModule):
|
||||
def __init__(self, config, output_dir,is_train=True):
|
||||
def __init__(self, config, output_dir, is_train=True):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.top_k = 3
|
||||
self.model = Text2SemanticDecoder(config=config, top_k=self.top_k)
|
||||
pretrained_s1=config.get("pretrained_s1")
|
||||
if(pretrained_s1 and is_train):
|
||||
pretrained_s1 = config.get("pretrained_s1")
|
||||
if pretrained_s1 and is_train:
|
||||
# print(self.load_state_dict(torch.load(pretrained_s1,map_location="cpu")["state_dict"]))
|
||||
print(self.load_state_dict(torch.load(pretrained_s1,map_location="cpu")["weight"]))
|
||||
print(
|
||||
self.load_state_dict(
|
||||
torch.load(pretrained_s1, map_location="cpu")["weight"]
|
||||
)
|
||||
)
|
||||
if is_train:
|
||||
self.automatic_optimization = False
|
||||
self.save_hyperparameters()
|
||||
self.eval_dir = output_dir / 'eval'
|
||||
self.eval_dir = output_dir / "eval"
|
||||
self.eval_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def training_step(self, batch: Dict, batch_idx: int):
|
||||
|
||||
opt = self.optimizers()
|
||||
scheduler = self.lr_schedulers()
|
||||
loss, acc = self.model.forward(
|
||||
batch['phoneme_ids'], batch['phoneme_ids_len'],
|
||||
batch['semantic_ids'], batch['semantic_ids_len'],
|
||||
batch['bert_feature'])
|
||||
batch["phoneme_ids"],
|
||||
batch["phoneme_ids_len"],
|
||||
batch["semantic_ids"],
|
||||
batch["semantic_ids_len"],
|
||||
batch["bert_feature"],
|
||||
)
|
||||
self.manual_backward(loss)
|
||||
if batch_idx > 0 and batch_idx % 4 == 0:
|
||||
opt.step()
|
||||
@@ -47,63 +54,67 @@ class Text2SemanticLightningModule(LightningModule):
|
||||
on_step=True,
|
||||
on_epoch=True,
|
||||
prog_bar=True,
|
||||
sync_dist=True)
|
||||
sync_dist=True,
|
||||
)
|
||||
self.log(
|
||||
"lr",
|
||||
scheduler.get_last_lr()[0],
|
||||
on_epoch=True,
|
||||
prog_bar=True,
|
||||
sync_dist=True)
|
||||
sync_dist=True,
|
||||
)
|
||||
self.log(
|
||||
f"top_{self.top_k}_acc",
|
||||
acc,
|
||||
on_step=True,
|
||||
on_epoch=True,
|
||||
prog_bar=True,
|
||||
sync_dist=True)
|
||||
sync_dist=True,
|
||||
)
|
||||
|
||||
def validation_step(self, batch: Dict, batch_idx: int):return
|
||||
# # get loss
|
||||
# loss, acc = self.model.forward(
|
||||
# batch['phoneme_ids'], batch['phoneme_ids_len'],
|
||||
# batch['semantic_ids'], batch['semantic_ids_len'],
|
||||
# batch['bert_feature']
|
||||
# )
|
||||
#
|
||||
# self.log(
|
||||
# "val_total_loss",
|
||||
# loss,
|
||||
# on_step=True,
|
||||
# on_epoch=True,
|
||||
# prog_bar=True,
|
||||
# sync_dist=True)
|
||||
# self.log(
|
||||
# f"val_top_{self.top_k}_acc",
|
||||
# acc,
|
||||
# on_step=True,
|
||||
# on_epoch=True,
|
||||
# prog_bar=True,
|
||||
# sync_dist=True)
|
||||
#
|
||||
# # get infer output
|
||||
# semantic_len = batch['semantic_ids'].size(1)
|
||||
# prompt_len = min(int(semantic_len * 0.5), 150)
|
||||
# prompt = batch['semantic_ids'][:, :prompt_len]
|
||||
# pred_semantic = self.model.infer(batch['phoneme_ids'],
|
||||
# batch['phoneme_ids_len'], prompt,
|
||||
# batch['bert_feature']
|
||||
# )
|
||||
# save_name = f'semantic_toks_{batch_idx}.pt'
|
||||
# save_path = os.path.join(self.eval_dir, save_name)
|
||||
# torch.save(pred_semantic.detach().cpu(), save_path)
|
||||
def validation_step(self, batch: Dict, batch_idx: int):
|
||||
return
|
||||
|
||||
# # get loss
|
||||
# loss, acc = self.model.forward(
|
||||
# batch['phoneme_ids'], batch['phoneme_ids_len'],
|
||||
# batch['semantic_ids'], batch['semantic_ids_len'],
|
||||
# batch['bert_feature']
|
||||
# )
|
||||
#
|
||||
# self.log(
|
||||
# "val_total_loss",
|
||||
# loss,
|
||||
# on_step=True,
|
||||
# on_epoch=True,
|
||||
# prog_bar=True,
|
||||
# sync_dist=True)
|
||||
# self.log(
|
||||
# f"val_top_{self.top_k}_acc",
|
||||
# acc,
|
||||
# on_step=True,
|
||||
# on_epoch=True,
|
||||
# prog_bar=True,
|
||||
# sync_dist=True)
|
||||
#
|
||||
# # get infer output
|
||||
# semantic_len = batch['semantic_ids'].size(1)
|
||||
# prompt_len = min(int(semantic_len * 0.5), 150)
|
||||
# prompt = batch['semantic_ids'][:, :prompt_len]
|
||||
# pred_semantic = self.model.infer(batch['phoneme_ids'],
|
||||
# batch['phoneme_ids_len'], prompt,
|
||||
# batch['bert_feature']
|
||||
# )
|
||||
# save_name = f'semantic_toks_{batch_idx}.pt'
|
||||
# save_path = os.path.join(self.eval_dir, save_name)
|
||||
# torch.save(pred_semantic.detach().cpu(), save_path)
|
||||
|
||||
def configure_optimizers(self):
|
||||
model_parameters = self.model.parameters()
|
||||
parameters_names = []
|
||||
parameters_names.append([
|
||||
name_param_pair[0]
|
||||
for name_param_pair in self.model.named_parameters()
|
||||
])
|
||||
parameters_names.append(
|
||||
[name_param_pair[0] for name_param_pair in self.model.named_parameters()]
|
||||
)
|
||||
lm_opt = ScaledAdam(
|
||||
model_parameters,
|
||||
lr=0.01,
|
||||
@@ -111,18 +122,19 @@ class Text2SemanticLightningModule(LightningModule):
|
||||
clipping_scale=2.0,
|
||||
parameters_names=parameters_names,
|
||||
show_dominant_parameters=False,
|
||||
clipping_update_period=1000, )
|
||||
clipping_update_period=1000,
|
||||
)
|
||||
|
||||
return {
|
||||
"optimizer": lm_opt,
|
||||
"lr_scheduler": {
|
||||
"scheduler":
|
||||
WarmupCosineLRSchedule(
|
||||
"scheduler": WarmupCosineLRSchedule(
|
||||
lm_opt,
|
||||
init_lr=self.config['optimizer']['lr_init'],
|
||||
peak_lr=self.config['optimizer']['lr'],
|
||||
end_lr=self.config['optimizer']['lr_end'],
|
||||
warmup_steps=self.config['optimizer']['warmup_steps'],
|
||||
total_steps=self.config['optimizer']['decay_steps'])
|
||||
}
|
||||
init_lr=self.config["optimizer"]["lr_init"],
|
||||
peak_lr=self.config["optimizer"]["lr"],
|
||||
end_lr=self.config["optimizer"]["lr_end"],
|
||||
warmup_steps=self.config["optimizer"]["warmup_steps"],
|
||||
total_steps=self.config["optimizer"]["decay_steps"],
|
||||
)
|
||||
},
|
||||
}
|
||||
|
||||
@@ -3,7 +3,12 @@ import torch
|
||||
from tqdm import tqdm
|
||||
|
||||
from AR.models.utils import make_pad_mask
|
||||
from AR.models.utils import topk_sampling,sample,logits_to_probs,multinomial_sample_one_no_sync
|
||||
from AR.models.utils import (
|
||||
topk_sampling,
|
||||
sample,
|
||||
logits_to_probs,
|
||||
multinomial_sample_one_no_sync,
|
||||
)
|
||||
from AR.modules.embedding import SinePositionalEmbedding
|
||||
from AR.modules.embedding import TokenEmbedding
|
||||
from AR.modules.transformer import LayerNorm
|
||||
@@ -22,35 +27,39 @@ default_config = {
|
||||
"p_dropout": 0.0,
|
||||
"vocab_size": 1024 + 1,
|
||||
"phoneme_vocab_size": 512,
|
||||
"EOS": 1024
|
||||
"EOS": 1024,
|
||||
}
|
||||
|
||||
|
||||
class Text2SemanticDecoder(nn.Module):
|
||||
def __init__(self, config, norm_first=False, top_k=3):
|
||||
super(Text2SemanticDecoder, self).__init__()
|
||||
self.model_dim = config['model']["hidden_dim"]
|
||||
self.embedding_dim = config['model']["embedding_dim"]
|
||||
self.num_head = config['model']["head"]
|
||||
self.num_layers = config['model']["n_layer"]
|
||||
self.model_dim = config["model"]["hidden_dim"]
|
||||
self.embedding_dim = config["model"]["embedding_dim"]
|
||||
self.num_head = config["model"]["head"]
|
||||
self.num_layers = config["model"]["n_layer"]
|
||||
self.norm_first = norm_first
|
||||
self.vocab_size = config['model']["vocab_size"]
|
||||
self.phoneme_vocab_size = config['model']["phoneme_vocab_size"]
|
||||
self.p_dropout = config['model']["dropout"]
|
||||
self.EOS = config['model']["EOS"]
|
||||
self.vocab_size = config["model"]["vocab_size"]
|
||||
self.phoneme_vocab_size = config["model"]["phoneme_vocab_size"]
|
||||
self.p_dropout = config["model"]["dropout"]
|
||||
self.EOS = config["model"]["EOS"]
|
||||
self.norm_first = norm_first
|
||||
assert self.EOS == self.vocab_size - 1
|
||||
# should be same as num of kmeans bin
|
||||
# assert self.EOS == 1024
|
||||
self.bert_proj = nn.Linear(1024, self.embedding_dim)
|
||||
self.ar_text_embedding = TokenEmbedding(
|
||||
self.embedding_dim, self.phoneme_vocab_size, self.p_dropout)
|
||||
self.embedding_dim, self.phoneme_vocab_size, self.p_dropout
|
||||
)
|
||||
self.ar_text_position = SinePositionalEmbedding(
|
||||
self.embedding_dim, dropout=0.1, scale=False, alpha=True)
|
||||
self.embedding_dim, dropout=0.1, scale=False, alpha=True
|
||||
)
|
||||
self.ar_audio_embedding = TokenEmbedding(
|
||||
self.embedding_dim, self.vocab_size, self.p_dropout)
|
||||
self.embedding_dim, self.vocab_size, self.p_dropout
|
||||
)
|
||||
self.ar_audio_position = SinePositionalEmbedding(
|
||||
self.embedding_dim, dropout=0.1, scale=False, alpha=True)
|
||||
self.embedding_dim, dropout=0.1, scale=False, alpha=True
|
||||
)
|
||||
|
||||
self.h = TransformerEncoder(
|
||||
TransformerEncoderLayer(
|
||||
@@ -59,28 +68,30 @@ class Text2SemanticDecoder(nn.Module):
|
||||
dim_feedforward=self.model_dim * 4,
|
||||
dropout=0.1,
|
||||
batch_first=True,
|
||||
norm_first=norm_first, ),
|
||||
norm_first=norm_first,
|
||||
),
|
||||
num_layers=self.num_layers,
|
||||
norm=LayerNorm(self.model_dim) if norm_first else None, )
|
||||
norm=LayerNorm(self.model_dim) if norm_first else None,
|
||||
)
|
||||
|
||||
self.ar_predict_layer = nn.Linear(
|
||||
self.model_dim, self.vocab_size, bias=False)
|
||||
self.loss_fct = nn.CrossEntropyLoss(reduction='sum')
|
||||
self.ar_predict_layer = nn.Linear(self.model_dim, self.vocab_size, bias=False)
|
||||
self.loss_fct = nn.CrossEntropyLoss(reduction="sum")
|
||||
|
||||
self.ar_accuracy_metric = MulticlassAccuracy(
|
||||
self.vocab_size,
|
||||
top_k=top_k,
|
||||
average="micro",
|
||||
multidim_average="global",
|
||||
ignore_index=self.EOS, )
|
||||
ignore_index=self.EOS,
|
||||
)
|
||||
|
||||
def forward(self, x, x_lens, y, y_lens, bert_feature):
|
||||
'''
|
||||
"""
|
||||
x: phoneme_ids
|
||||
y: semantic_ids
|
||||
'''
|
||||
"""
|
||||
x = self.ar_text_embedding(x)
|
||||
x = x + self.bert_proj(bert_feature.transpose(1,2))
|
||||
x = x + self.bert_proj(bert_feature.transpose(1, 2))
|
||||
x = self.ar_text_position(x)
|
||||
x_mask = make_pad_mask(x_lens)
|
||||
|
||||
@@ -102,18 +113,23 @@ class Text2SemanticDecoder(nn.Module):
|
||||
x_attn_mask = F.pad(
|
||||
torch.zeros((x_len, x_len), dtype=torch.bool, device=x.device),
|
||||
(0, y_len),
|
||||
value=True, )
|
||||
value=True,
|
||||
)
|
||||
y_attn_mask = F.pad(
|
||||
torch.triu(
|
||||
torch.ones(y_len, y_len, dtype=torch.bool, device=x.device),
|
||||
diagonal=1, ),
|
||||
diagonal=1,
|
||||
),
|
||||
(x_len, 0),
|
||||
value=False, )
|
||||
value=False,
|
||||
)
|
||||
xy_attn_mask = torch.concat([x_attn_mask, y_attn_mask], dim=0)
|
||||
bsz, src_len = x.shape[0], x_len + y_len
|
||||
_xy_padding_mask = (ar_xy_padding_mask.view(bsz, 1, 1, src_len)
|
||||
.expand(-1, self.num_head, -1, -1)
|
||||
.reshape(bsz * self.num_head, 1, src_len))
|
||||
_xy_padding_mask = (
|
||||
ar_xy_padding_mask.view(bsz, 1, 1, src_len)
|
||||
.expand(-1, self.num_head, -1, -1)
|
||||
.reshape(bsz * self.num_head, 1, src_len)
|
||||
)
|
||||
xy_attn_mask = xy_attn_mask.logical_or(_xy_padding_mask)
|
||||
new_attn_mask = torch.zeros_like(xy_attn_mask, dtype=x.dtype)
|
||||
new_attn_mask.masked_fill_(xy_attn_mask, float("-inf"))
|
||||
@@ -122,26 +138,28 @@ class Text2SemanticDecoder(nn.Module):
|
||||
xy_pos = torch.concat([x, y_pos], dim=1)
|
||||
xy_dec, _ = self.h(
|
||||
(xy_pos, None),
|
||||
mask=xy_attn_mask, )
|
||||
mask=xy_attn_mask,
|
||||
)
|
||||
logits = self.ar_predict_layer(xy_dec[:, x_len:]).permute(0, 2, 1)
|
||||
# loss
|
||||
# from feiteng: 每次 duration 越多, 梯度更新也应该更多, 所以用 sum
|
||||
loss = F.cross_entropy(logits, targets, reduction='sum')
|
||||
loss = F.cross_entropy(logits, targets, reduction="sum")
|
||||
acc = self.ar_accuracy_metric(logits.detach(), targets).item()
|
||||
return loss, acc
|
||||
|
||||
# 需要看下这个函数和 forward 的区别以及没有 semantic 的时候 prompts 输入什么
|
||||
def infer(self,
|
||||
x,
|
||||
x_lens,
|
||||
prompts,
|
||||
bert_feature,
|
||||
top_k: int=-100,
|
||||
early_stop_num: int=-1,
|
||||
temperature: float=1.0):
|
||||
|
||||
def infer(
|
||||
self,
|
||||
x,
|
||||
x_lens,
|
||||
prompts,
|
||||
bert_feature,
|
||||
top_k: int = -100,
|
||||
early_stop_num: int = -1,
|
||||
temperature: float = 1.0,
|
||||
):
|
||||
x = self.ar_text_embedding(x)
|
||||
x = x + self.bert_proj(bert_feature.transpose(1,2))
|
||||
x = x + self.bert_proj(bert_feature.transpose(1, 2))
|
||||
x = self.ar_text_position(x)
|
||||
|
||||
# AR Decoder
|
||||
@@ -159,35 +177,37 @@ class Text2SemanticDecoder(nn.Module):
|
||||
x_attn_mask_pad = F.pad(
|
||||
x_attn_mask,
|
||||
(0, y_len),
|
||||
value=True, )
|
||||
value=True,
|
||||
)
|
||||
y_attn_mask = F.pad(
|
||||
torch.triu(
|
||||
torch.ones(y_len, y_len, dtype=torch.bool), diagonal=1),
|
||||
torch.triu(torch.ones(y_len, y_len, dtype=torch.bool), diagonal=1),
|
||||
(x_len, 0),
|
||||
value=False, )
|
||||
xy_attn_mask = torch.concat(
|
||||
[x_attn_mask_pad, y_attn_mask], dim=0).to(y.device)
|
||||
value=False,
|
||||
)
|
||||
xy_attn_mask = torch.concat([x_attn_mask_pad, y_attn_mask], dim=0).to(
|
||||
y.device
|
||||
)
|
||||
|
||||
xy_dec, _ = self.h(
|
||||
(xy_pos, None),
|
||||
mask=xy_attn_mask, )
|
||||
mask=xy_attn_mask,
|
||||
)
|
||||
logits = self.ar_predict_layer(xy_dec[:, -1])
|
||||
samples = topk_sampling(
|
||||
logits, top_k=top_k, top_p=1.0, temperature=temperature)
|
||||
logits, top_k=top_k, top_p=1.0, temperature=temperature
|
||||
)
|
||||
|
||||
if early_stop_num != -1 and (y.shape[1] - prefix_len
|
||||
) > early_stop_num:
|
||||
if early_stop_num != -1 and (y.shape[1] - prefix_len) > early_stop_num:
|
||||
print("use early stop num:", early_stop_num)
|
||||
stop = True
|
||||
|
||||
if torch.argmax(
|
||||
logits, dim=-1)[0] == self.EOS or samples[0, 0] == self.EOS:
|
||||
if torch.argmax(logits, dim=-1)[0] == self.EOS or samples[0, 0] == self.EOS:
|
||||
# print(torch.argmax(logits, dim=-1)[0] == self.EOS, samples[0, 0] == self.EOS)
|
||||
stop = True
|
||||
if stop:
|
||||
if prompts.shape[1] == y.shape[1]:
|
||||
y = torch.concat([y, torch.zeros_like(samples)], dim=1)
|
||||
print('bad zero prediction')
|
||||
print("bad zero prediction")
|
||||
print(f"T2S Decoding EOS [{prefix_len} -> {y.shape[1]}]")
|
||||
break
|
||||
# 本次生成的 semantic_ids 和之前的 y 构成新的 y
|
||||
@@ -198,23 +218,24 @@ class Text2SemanticDecoder(nn.Module):
|
||||
return y
|
||||
|
||||
def pad_y_eos(self, y, y_mask_int, eos_id):
|
||||
targets = F.pad(
|
||||
y, (0, 1), value=0) + eos_id * F.pad(
|
||||
y_mask_int, (0, 1), value=1)
|
||||
targets = F.pad(y, (0, 1), value=0) + eos_id * F.pad(
|
||||
y_mask_int, (0, 1), value=1
|
||||
)
|
||||
# 错位
|
||||
return targets[:, :-1], targets[:, 1:]
|
||||
|
||||
def infer_panel(self,
|
||||
x,#####全部文本token
|
||||
x_lens,
|
||||
prompts,####参考音频token
|
||||
bert_feature,
|
||||
top_k: int=-100,
|
||||
early_stop_num: int=-1,
|
||||
temperature: float=1.0):
|
||||
|
||||
def infer_panel(
|
||||
self,
|
||||
x, #####全部文本token
|
||||
x_lens,
|
||||
prompts, ####参考音频token
|
||||
bert_feature,
|
||||
top_k: int = -100,
|
||||
early_stop_num: int = -1,
|
||||
temperature: float = 1.0,
|
||||
):
|
||||
x = self.ar_text_embedding(x)
|
||||
x = x + self.bert_proj(bert_feature.transpose(1,2))
|
||||
x = x + self.bert_proj(bert_feature.transpose(1, 2))
|
||||
x = self.ar_text_position(x)
|
||||
|
||||
# AR Decoder
|
||||
@@ -224,75 +245,81 @@ class Text2SemanticDecoder(nn.Module):
|
||||
x_attn_mask = torch.zeros((x_len, x_len), dtype=torch.bool)
|
||||
stop = False
|
||||
# print(1111111,self.num_layers)
|
||||
cache={
|
||||
"all_stage":self.num_layers,
|
||||
"k":[None]*self.num_layers,###根据配置自己手写
|
||||
"v":[None]*self.num_layers,
|
||||
cache = {
|
||||
"all_stage": self.num_layers,
|
||||
"k": [None] * self.num_layers, ###根据配置自己手写
|
||||
"v": [None] * self.num_layers,
|
||||
# "xy_pos":None,##y_pos位置编码每次都不一样的没法缓存，每次都要重新拼xy_pos.主要还是写法原因，其实是可以历史统一一样的，但也没啥计算量就不管了
|
||||
"y_emb":None,##只需要对最新的samples求emb，再拼历史的就行
|
||||
"y_emb": None, ##只需要对最新的samples求emb，再拼历史的就行
|
||||
# "logits":None,###原版就已经只对结尾求再拼接了，不用管
|
||||
# "xy_dec":None,###不需要，本来只需要最后一个做logits
|
||||
"first_infer":1,
|
||||
"stage":0
|
||||
"first_infer": 1,
|
||||
"stage": 0,
|
||||
}
|
||||
for idx in tqdm(range(1500)):
|
||||
if(cache["first_infer"]==1):
|
||||
if cache["first_infer"] == 1:
|
||||
y_emb = self.ar_audio_embedding(y)
|
||||
else:
|
||||
y_emb = torch.cat([cache["y_emb"],self.ar_audio_embedding(y[:,-1:])],1)
|
||||
cache["y_emb"]=y_emb
|
||||
y_emb = torch.cat(
|
||||
[cache["y_emb"], self.ar_audio_embedding(y[:, -1:])], 1
|
||||
)
|
||||
cache["y_emb"] = y_emb
|
||||
y_pos = self.ar_audio_position(y_emb)
|
||||
# x 和逐渐增长的 y 一起输入给模型
|
||||
if(cache["first_infer"]==1):
|
||||
if cache["first_infer"] == 1:
|
||||
xy_pos = torch.concat([x, y_pos], dim=1)
|
||||
else:
|
||||
xy_pos=y_pos[:,-1:]
|
||||
xy_pos = y_pos[:, -1:]
|
||||
y_len = y_pos.shape[1]
|
||||
###以下3个不做缓存
|
||||
if (cache["first_infer"] == 1):
|
||||
if cache["first_infer"] == 1:
|
||||
x_attn_mask_pad = F.pad(
|
||||
x_attn_mask,
|
||||
(0, y_len),###xx的纯0扩展到xx纯0+xy纯1，(x,x+y)
|
||||
value=True, )
|
||||
y_attn_mask = F.pad(###yy的右上1扩展到左边xy的0,(y,x+y)
|
||||
torch.triu(
|
||||
torch.ones(y_len, y_len, dtype=torch.bool), diagonal=1),
|
||||
x_attn_mask,
|
||||
(0, y_len), ###xx的纯0扩展到xx纯0+xy纯1，(x,x+y)
|
||||
value=True,
|
||||
)
|
||||
y_attn_mask = F.pad( ###yy的右上1扩展到左边xy的0,(y,x+y)
|
||||
torch.triu(torch.ones(y_len, y_len, dtype=torch.bool), diagonal=1),
|
||||
(x_len, 0),
|
||||
value=False, )
|
||||
xy_attn_mask = torch.concat(
|
||||
[x_attn_mask_pad, y_attn_mask], dim=0).to(y.device)
|
||||
value=False,
|
||||
)
|
||||
xy_attn_mask = torch.concat([x_attn_mask_pad, y_attn_mask], dim=0).to(
|
||||
y.device
|
||||
)
|
||||
else:
|
||||
###最右边一列（是错的）
|
||||
# xy_attn_mask=torch.ones((1, x_len+y_len), dtype=torch.bool,device=xy_pos.device)
|
||||
# xy_attn_mask[:,-1]=False
|
||||
###最下面一行（是对的）
|
||||
xy_attn_mask = torch.zeros((1, x_len + y_len), dtype=torch.bool, device=xy_pos.device)
|
||||
xy_attn_mask = torch.zeros(
|
||||
(1, x_len + y_len), dtype=torch.bool, device=xy_pos.device
|
||||
)
|
||||
# pdb.set_trace()
|
||||
###缓存重头戏
|
||||
# print(1111,xy_pos.shape,xy_attn_mask.shape,x_len,y_len)
|
||||
xy_dec, _ = self.h(
|
||||
(xy_pos, None),
|
||||
mask=xy_attn_mask,cache=cache )
|
||||
logits = self.ar_predict_layer(xy_dec[:, -1])##不用改，如果用了cache的默认就是只有一帧，取最后一帧一样的
|
||||
xy_dec, _ = self.h((xy_pos, None), mask=xy_attn_mask, cache=cache)
|
||||
logits = self.ar_predict_layer(
|
||||
xy_dec[:, -1]
|
||||
) ##不用改，如果用了cache的默认就是只有一帧，取最后一帧一样的
|
||||
# samples = topk_sampling(logits, top_k=top_k, top_p=1.0, temperature=temperature)
|
||||
samples = sample(logits[0], y, top_k=top_k, top_p=1.0, repetition_penalty=1.35)[0].unsqueeze(0)
|
||||
if early_stop_num != -1 and (y.shape[1] - prefix_len
|
||||
) > early_stop_num:
|
||||
samples = sample(
|
||||
logits[0], y, top_k=top_k, top_p=1.0, repetition_penalty=1.35
|
||||
)[0].unsqueeze(0)
|
||||
if early_stop_num != -1 and (y.shape[1] - prefix_len) > early_stop_num:
|
||||
print("use early stop num:", early_stop_num)
|
||||
stop = True
|
||||
|
||||
if torch.argmax(
|
||||
logits, dim=-1)[0] == self.EOS or samples[0, 0] == self.EOS:
|
||||
if torch.argmax(logits, dim=-1)[0] == self.EOS or samples[0, 0] == self.EOS:
|
||||
# print(torch.argmax(logits, dim=-1)[0] == self.EOS, samples[0, 0] == self.EOS)
|
||||
stop = True
|
||||
if stop:
|
||||
if prompts.shape[1] == y.shape[1]:
|
||||
y = torch.concat([y, torch.zeros_like(samples)], dim=1)
|
||||
print('bad zero prediction')
|
||||
print("bad zero prediction")
|
||||
print(f"T2S Decoding EOS [{prefix_len} -> {y.shape[1]}]")
|
||||
break
|
||||
# 本次生成的 semantic_ids 和之前的 y 构成新的 y
|
||||
# print(samples.shape)#[1,1]#第一个1是bs
|
||||
y = torch.concat([y, samples], dim=1)
|
||||
cache["first_infer"]=0
|
||||
return y,idx
|
||||
cache["first_infer"] = 0
|
||||
return y, idx
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def sequence_mask(length, max_length=None):
|
||||
if max_length is None:
|
||||
max_length = length.max()
|
||||
@@ -9,7 +10,7 @@ def sequence_mask(length, max_length=None):
|
||||
return x.unsqueeze(0) < length.unsqueeze(1)
|
||||
|
||||
|
||||
def make_pad_mask(lengths: torch.Tensor, max_len: int=0) -> torch.Tensor:
|
||||
def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
lengths:
|
||||
@@ -38,11 +39,9 @@ def make_pad_mask(lengths: torch.Tensor, max_len: int=0) -> torch.Tensor:
|
||||
|
||||
|
||||
# https://github.com/microsoft/unilm/blob/master/xtune/src/transformers/modeling_utils.py
|
||||
def top_k_top_p_filtering(logits,
|
||||
top_k=0,
|
||||
top_p=1.0,
|
||||
filter_value=-float("Inf"),
|
||||
min_tokens_to_keep=1):
|
||||
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)
|
||||
@@ -53,16 +52,14 @@ def top_k_top_p_filtering(logits,
|
||||
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
|
||||
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)
|
||||
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
|
||||
@@ -70,13 +67,13 @@ def top_k_top_p_filtering(logits,
|
||||
# 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[..., 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)
|
||||
1, sorted_indices, sorted_indices_to_remove
|
||||
)
|
||||
logits[indices_to_remove] = filter_value
|
||||
return logits
|
||||
|
||||
@@ -100,6 +97,8 @@ def topk_sampling(logits, top_k=10, top_p=1.0, temperature=1.0):
|
||||
|
||||
|
||||
from typing import Optional, Tuple
|
||||
|
||||
|
||||
def multinomial_sample_one_no_sync(
|
||||
probs_sort,
|
||||
): # Does multinomial sampling without a cuda synchronization
|
||||
@@ -115,7 +114,7 @@ def logits_to_probs(
|
||||
top_p: Optional[int] = None,
|
||||
repetition_penalty: float = 1.0,
|
||||
):
|
||||
previous_tokens=previous_tokens.squeeze()
|
||||
previous_tokens = previous_tokens.squeeze()
|
||||
# print(logits.shape,previous_tokens.shape)
|
||||
# pdb.set_trace()
|
||||
if previous_tokens is not None and repetition_penalty != 1.0:
|
||||
@@ -159,4 +158,3 @@ def sample(
|
||||
)
|
||||
idx_next = multinomial_sample_one_no_sync(probs)
|
||||
return idx_next, probs
|
||||
|
||||
|
||||
Reference in New Issue
Block a user