304 lines
7.9 KiB
Python
304 lines
7.9 KiB
Python
import os
|
|
import json
|
|
import torch
|
|
import torch.nn as nn
|
|
from pathlib import Path
|
|
from dataclasses import dataclass
|
|
from typing import Union
|
|
from torch.optim.lr_scheduler import LambdaLR
|
|
from transformers import AutoTokenizer
|
|
from fireredtts2.llm.llm import Model, ModelArgs
|
|
|
|
|
|
@dataclass
|
|
class Segment:
|
|
speaker: str
|
|
text: str
|
|
audio: torch.Tensor
|
|
|
|
|
|
class WarmupDecayLR(LambdaLR):
|
|
"""
|
|
Learning rate scheduler with a linear warmup and specificable decay.
|
|
"""
|
|
|
|
def __init__(
|
|
self, optimizer, warmup_steps: int, total_steps: int, decay_type: str = "linear"
|
|
):
|
|
self.warmup_steps = warmup_steps
|
|
self.total_steps = total_steps
|
|
self.decay_type = decay_type
|
|
super().__init__(optimizer, self.lr_lambda, last_epoch=-1)
|
|
|
|
def lr_lambda(self, step: int) -> float:
|
|
if step < self.warmup_steps:
|
|
return step / self.warmup_steps
|
|
else:
|
|
if self.decay_type == "linear":
|
|
return (self.total_steps - step) / (
|
|
self.total_steps - self.warmup_steps
|
|
)
|
|
elif self.decay_type == "constant":
|
|
return 1.0
|
|
elif self.decay_type == "exponential":
|
|
return 0.1 ** (
|
|
(step - self.warmup_steps) / (self.total_steps - self.warmup_steps)
|
|
)
|
|
elif self.decay_type == "cosine":
|
|
return 0.5 * (
|
|
1
|
|
+ torch.cos(
|
|
torch.pi
|
|
* torch.tensor(
|
|
(step - self.warmup_steps)
|
|
/ (self.total_steps - self.warmup_steps)
|
|
)
|
|
)
|
|
)
|
|
else:
|
|
raise ValueError(f"Invalid decay type: {self.decay_type}")
|
|
|
|
|
|
additional_special_tokens = [
|
|
"<|text_start|>",
|
|
"<|text_end|>",
|
|
"[S1]",
|
|
"[S2]",
|
|
"[S3]",
|
|
"[S4]",
|
|
"[S5]",
|
|
"[S6]",
|
|
"[S7]",
|
|
"[S8]",
|
|
"[S9]",
|
|
"[S10]",
|
|
"[S11]",
|
|
"[S12]",
|
|
"[S13]",
|
|
"[S14]",
|
|
"[S15]",
|
|
"[S16]",
|
|
"[S17]",
|
|
"[S18]",
|
|
"[S19]",
|
|
"[S20]",
|
|
"[S21]",
|
|
"[S22]",
|
|
"[S23]",
|
|
"[S24]",
|
|
"[S25]",
|
|
"[S26]",
|
|
"[S27]",
|
|
"[S28]",
|
|
"[S29]",
|
|
"[S30]",
|
|
"[S31]",
|
|
"[S32]",
|
|
"[S33]",
|
|
"[S34]",
|
|
"[S35]",
|
|
"[S36]",
|
|
"[S37]",
|
|
"[S38]",
|
|
"[S39]",
|
|
"[S40]",
|
|
"[S_PODCAST_1]",
|
|
"[S_PODCAST_2]",
|
|
"[S_PODCAST_3]",
|
|
"[S_PODCAST_4]",
|
|
"[S_PODCAST_5]",
|
|
"[S_PODCAST_6]",
|
|
"[S_PODCAST_7]",
|
|
"[S_PODCAST_8]",
|
|
"[S_PODCAST_9]",
|
|
"[S_PODCAST_10]",
|
|
"[S_DIALOG_1]",
|
|
"[S_DIALOG_2]",
|
|
"[S_DIALOG_3]",
|
|
"[S_DIALOG_4]",
|
|
"[S_DIALOG_5]",
|
|
"[S_DIALOG_6]",
|
|
"[S_DIALOG_7]",
|
|
"[S_DIALOG_8]",
|
|
"[S_DIALOG_9]",
|
|
"[S_DIALOG_10]",
|
|
"<|emotion_neutral|>",
|
|
"<|emotion_happy|>",
|
|
"<|emotion_sad|>",
|
|
"<|emotion_concern|>",
|
|
"<|emotion_confuse|>",
|
|
"<|emotion_angry|>",
|
|
"<|emotion_surprise|>",
|
|
"<|emotion_disgust|>",
|
|
"<|emotion_nervous|>",
|
|
"<|emotion_apology|>",
|
|
"<|emotion_understand|>",
|
|
"<|emotion_fear|>",
|
|
"<|emotion_comfort|>",
|
|
"<|emotion_shy|>",
|
|
"<|emotion_serious|>",
|
|
"<|emotion_extra1|>",
|
|
"<|emotion_extra2|>",
|
|
"<|emotion_extra3|>",
|
|
"<|emotion_extra4|>",
|
|
"<|emotion_extra5|>",
|
|
"<|emotion_extra6|>",
|
|
"<|emotion_extra7|>",
|
|
"<|emotion_extra8|>",
|
|
"<|emotion_extra9|>",
|
|
"<|emotion_extra10|>",
|
|
"<|breath|>",
|
|
"<|humph|>",
|
|
"<|laugh_heng|>",
|
|
"<|hissing|>",
|
|
"<|sniff|>",
|
|
"<|laugh_he|>",
|
|
"<|sigh|>",
|
|
"<|laugh|>",
|
|
"<|laugh_ha|>",
|
|
"<|quick_breath|>",
|
|
"<|laugh_hei|>",
|
|
"<|laugh_speak|>",
|
|
"<|/laugh_speak|>",
|
|
"<|cry|>",
|
|
"<|choking|>",
|
|
"<|cry_speak|>",
|
|
"<|/cry_speak|>",
|
|
"<|slurp|>",
|
|
"<|clucking|>",
|
|
"<|yawning|>",
|
|
"<|cough|>",
|
|
"<|smack|>",
|
|
"<|hem|>",
|
|
"<|stretch|>",
|
|
"<|sneeze|>",
|
|
"<|paralinguistic_extra1|>",
|
|
"<|paralinguistic_extra2|>",
|
|
"<|paralinguistic_extra3|>",
|
|
"<|paralinguistic_extra4|>",
|
|
"<|paralinguistic_extra5|>",
|
|
"<|paralinguistic_extra6|>",
|
|
"<|paralinguistic_extra7|>",
|
|
"<|paralinguistic_extra8|>",
|
|
"<|paralinguistic_extra10|>",
|
|
"<|paralinguistic_extra11|>",
|
|
"<|paralinguistic_extra12|>",
|
|
"<|paralinguistic_extra13|>",
|
|
]
|
|
|
|
|
|
def load_custom_tokenizer(qwen2_tokenizer_path: str):
|
|
tok = AutoTokenizer.from_pretrained(qwen2_tokenizer_path)
|
|
special_tokens_dict = {
|
|
"additional_special_tokens": additional_special_tokens,
|
|
}
|
|
tok.add_special_tokens(special_tokens_dict)
|
|
return tok
|
|
|
|
|
|
def init_weights(model: nn.Module):
|
|
"""
|
|
Initialize the weights of the model.
|
|
- Xavier uniform initialization for linear layers
|
|
- Normal initialization for embeddings
|
|
- Xavier uniform initialization for parameters
|
|
"""
|
|
|
|
def _init_weights(m):
|
|
if isinstance(m, nn.Linear):
|
|
nn.init.xavier_uniform_(m.weight)
|
|
if m.bias is not None:
|
|
nn.init.zeros_(m.bias)
|
|
elif isinstance(m, nn.Embedding):
|
|
nn.init.normal_(m.weight, mean=0.0, std=0.02)
|
|
elif isinstance(m, nn.Parameter):
|
|
nn.init.xavier_uniform_(m.data)
|
|
|
|
model.apply(_init_weights)
|
|
|
|
# Special handling for audio_head because it's nn.Parameter directly
|
|
nn.init.xavier_uniform_(model.audio_head)
|
|
|
|
return model
|
|
|
|
|
|
def load_llm_model(
|
|
configs,
|
|
checkpoint_path: Union[str, Path] = None,
|
|
device: Union[str, torch.device] = "cuda",
|
|
) -> Model:
|
|
"""Load model, add forward method, and move to device.
|
|
|
|
Args:
|
|
model_name_or_checkpoint_path: Name or path of pretrained model or checkpoint.
|
|
device: Device to move the model to.
|
|
decoder_loss_weight: Decoder loss weight.
|
|
"""
|
|
|
|
model_arg = ModelArgs(
|
|
backbone_flavor=configs["llm_models"]["backbone_flavor"],
|
|
decoder_flavor=configs["llm_models"]["decoder_flavor"],
|
|
text_vocab_size=configs["llm_models"]["text_vocab_size"],
|
|
audio_vocab_size=configs["llm_models"]["audio_vocab_size"],
|
|
audio_num_codebooks=configs["llm_models"]["audio_num_codebooks"],
|
|
decoder_loss_weight=configs["llm_models"]["decoder_loss_weight"],
|
|
use_text_loss=True,
|
|
)
|
|
model = Model(model_arg)
|
|
|
|
if checkpoint_path and os.path.exists(checkpoint_path):
|
|
state_dict = torch.load(
|
|
checkpoint_path, map_location="cpu", weights_only=False
|
|
)["model"]
|
|
model.load_state_dict(state_dict)
|
|
else:
|
|
model = init_weights(model)
|
|
|
|
model = model.to(device=device)
|
|
return model
|
|
|
|
|
|
def summarize(
|
|
writer,
|
|
global_step,
|
|
scalars={},
|
|
histograms={},
|
|
images={},
|
|
audios={},
|
|
audio_sampling_rate=22050,
|
|
):
|
|
for k, v in scalars.items():
|
|
writer.add_scalar(k, v, global_step)
|
|
for k, v in histograms.items():
|
|
writer.add_histogram(k, v, global_step)
|
|
for k, v in images.items():
|
|
writer.add_image(k, v, global_step, dataformats="HWC")
|
|
for k, v in audios.items():
|
|
writer.add_audio(k, v, global_step, audio_sampling_rate)
|
|
|
|
|
|
def get_grad_norm(model):
|
|
total_norm = 0
|
|
num = 0
|
|
for name, p in model.named_parameters():
|
|
try:
|
|
param_norm = p.grad.data.norm(2)
|
|
total_norm += param_norm.item() ** 2
|
|
num += 1
|
|
except:
|
|
print(name)
|
|
total_norm = total_norm ** (1.0 / 2)
|
|
total_norm = total_norm / num
|
|
return total_norm
|
|
|
|
|
|
def read_jsonl(path):
|
|
path = os.path.expanduser(path)
|
|
with open(path, "r") as f:
|
|
json_str = f.read()
|
|
data_list = []
|
|
for line in json_str.splitlines():
|
|
data = json.loads(line)
|
|
data_list.append(data)
|
|
return data_list
|