import logging import os import torch import torch.nn as nn from funcineforge.models.utils.llm_decoding import LLMDecoder from funcineforge.utils.device_funcs import to_device import numpy as np from funcineforge.models.utils import dtype_map from funcineforge.models import FunCineForgeSpecAug from transformers import AutoModelForCausalLM import pickle class FunCineForgeLM(nn.Module): def __init__( self, llm: str = None, llm_conf: dict = None, input_size: int = 80, length_normalized_loss: bool = False, **kwargs, ): super().__init__() # llm self.llm_conf = llm_conf self.llm = None init_param_path = llm_conf.get("init_param_path", "") llm_load_kwargs = llm_conf.get("load_kwargs", {}) self.sample_rate = kwargs.get("sample_rate", 24000) self.token_rate = kwargs.get("token_rate", 25) if kwargs.get("infer_lora_merged", False): llm_conf["use_qlora"] = False llm_conf["use_lora"] = False kwargs["infer_use_lora"] = False model = AutoModelForCausalLM.from_pretrained( init_param_path, load_in_8bit=None, device_map=None, use_cache=None, **llm_load_kwargs, ) freeze = llm_conf.get("freeze", True) if freeze: for name, param in model.named_parameters(): param.requires_grad = False model.eval() logging.info(f"use_lora: {llm_conf.get('use_lora', False)}, use_qlora: {llm_conf.get('use_qlora', False)}, infer_use_lora: {kwargs.get('infer_use_lora',False)}, infer_lora_merged: {kwargs.get('infer_lora_merged',False)}") if llm_conf.get("activation_checkpoint", False): model.gradient_checkpointing_enable() self.llm_dtype = llm_conf.get("llm_dtype", "fp32") self.llm = model.to(dtype_map[self.llm_dtype]) llm_dim = model.get_input_embeddings().weight.shape[-1] if (not llm_conf.get("use_lora", False)) and (not kwargs.get("infer_use_lora",False)): del self.llm.lm_head self.codec_unit = kwargs.get("codec_unit", 6761) self.timespk_unit = kwargs.get("timespk_unit", 1550) self.codec_embed = nn.Embedding(self.codec_unit, llm_dim, 0) self.timespk_embed = nn.Embedding(self.timespk_unit, llm_dim, 0) self.codec_head = nn.Linear(llm_dim, self.codec_unit, bias=False) self.face_size = kwargs.get("face_size", 512) self.face_linear = nn.Linear(self.face_size, llm_dim) self.length_normalized_loss = length_normalized_loss self.ignore_id = kwargs.get("ignore_id", -100) specaug = kwargs.get("specaug", None) specaug_conf = kwargs.get("specaug_conf", {}) if specaug is not None: specaug = FunCineForgeSpecAug(**specaug_conf) self.specaug = specaug rank = int(os.environ.get("RANK", 0)) logging.info(f"rank: {rank}, model is builded.") def insert_face_embeddings( self, inputs_embeds, face_emb, attention_mask, labels_ids, codec_len, insert_pos, device ): """ 将face_emb插入到inputs_embeds中的指定位置, 同步更新attention_mask和labels_ids Args: inputs_embeds: (batch_size, token_num, dims) 输入embedding face_emb: (batch_size, max_face_len, dims) 面部embedding attention_mask: (batch_size, token_num) 注意力mask labels_ids: (batch_size, token_num) 标签ID codec_len: (batch_size,) 每个样本的实际face_emb长度 insert_pos: int 插入位置, SOS token之后 device Returns: padded_inputs_embeds: 插入face_emb并padding后的inputs_embeds padded_attention_mask: 更新后的attention_mask padded_labels: 更新后的labels_ids """ batch_size, token_num, dims = inputs_embeds.shape max_face_len = face_emb.size(1) # 预计算新序列的最大长度 new_max_length = token_num + max_face_len # 预分配输出张量 padded_inputs_embeds = torch.zeros(batch_size, new_max_length, dims, device=device) padded_attention_mask = torch.zeros(batch_size, new_max_length, device=device, dtype=attention_mask.dtype) padded_labels = torch.full((batch_size, new_max_length), self.ignore_id, device=device, dtype=labels_ids.dtype) for i in range(batch_size): current_face_len = codec_len[i].item() # 直接填充，避免中间拼接 padded_inputs_embeds[i, :insert_pos] = inputs_embeds[i, :insert_pos] padded_inputs_embeds[i, insert_pos:insert_pos+current_face_len] = face_emb[i, :current_face_len] padded_inputs_embeds[i, insert_pos+current_face_len:token_num+current_face_len] = inputs_embeds[i, insert_pos:] # 同样处理mask和labels padded_attention_mask[i, :insert_pos] = attention_mask[i, :insert_pos] padded_attention_mask[i, insert_pos:insert_pos+current_face_len] = 1 padded_attention_mask[i, insert_pos+current_face_len:token_num+current_face_len] = attention_mask[i, insert_pos:] padded_labels[i, :insert_pos] = labels_ids[i, :insert_pos] padded_labels[i, insert_pos:insert_pos+current_face_len] = self.ignore_id padded_labels[i, insert_pos+current_face_len:token_num+current_face_len] = labels_ids[i, insert_pos:] return padded_inputs_embeds, padded_attention_mask, padded_labels def load_data(self, contents: dict, **kwargs): lm_use_prompt = kwargs.get("lm_use_prompt", True) tokenizer = kwargs.get("tokenizer") # text + clue text = contents["text"] clue = "<|startofclue|>" + contents["clue"] + "<|endofclue|>" if lm_use_prompt: text = clue + text text_ids = tokenizer.encode(text) text_len = len(text_ids) # timespk_ids timespk_ids = contents["timespk_ids"].tolist() type_id = contents["type_id"] # sequence sequence = [ kwargs['dataset_conf']["sos"], *text_ids, type_id, *timespk_ids, kwargs['dataset_conf']["turn_of_speech"] ] input_ids = torch.tensor(sequence, dtype=torch.int64) # flag tensors text_flag = torch.zeros(len(sequence), dtype=torch.float32) timespk_flag = torch.zeros(len(sequence), dtype=torch.float32) codec_flag = torch.zeros(len(sequence), dtype=torch.float32) text_flag[1: text_len+1] = 1 timespk_flag[text_len+1: -1] = 1 codec_flag = 1 - text_flag - timespk_flag # face embs speech_len = contents["speech_len"] face_embs = torch.zeros((speech_len, self.face_size), dtype=torch.float32) face_path = contents.get("face") with open(face_path, 'rb') as f: stat_obj = pickle.load(f) embeddings = stat_obj['embeddings'] faceI = stat_obj['faceI'] for emb, frameI in zip(embeddings, faceI): fi = int(frameI) if 0 <= fi < speech_len: end = min(fi + 5, speech_len) face_embs[fi:end] = torch.from_numpy(emb).expand(end - fi, -1) # batch dimension input_ids = input_ids[None, :] text_flag = text_flag[None, :] timespk_flag = timespk_flag[None, :] codec_flag = codec_flag[None, :] face_embs = face_embs[None, :, :] output = { "input_ids": input_ids, "face_embs": face_embs, "text_flag": text_flag > 0, "timespk_flag": timespk_flag > 0, "codec_flag": codec_flag > 0, "prompt_codec": None, # you can add prompt codec here if needed } return output def inference_prepare(self, data_in, **kwargs): if kwargs.get("batch_size", 1) > 1: raise NotImplementedError("batch decoding is not implemented") output = self.load_data(data_in[0], **kwargs) batch = to_device(output, kwargs["device"]) input_ids = batch["input_ids"] input_ids = input_ids * (input_ids > 0) text_flag = batch["text_flag"] timespk_flag = batch["timespk_flag"] codec_flag = batch["codec_flag"] face_embs = batch["face_embs"] if (kwargs.get("use_qlora",False) or kwargs.get("infer_use_lora",False)) and (not kwargs.get("infer_lora_merged",False)): text_embeds = self.llm.base_model.model.model.get_input_embeddings()(input_ids * text_flag) * text_flag.unsqueeze(-1) else: text_embeds = self.llm.model.get_input_embeddings()(input_ids * text_flag) * text_flag.unsqueeze(-1) timespk_embeds = self.timespk_embed(input_ids * timespk_flag) * timespk_flag.unsqueeze(-1) codec_embs = self.codec_embed(input_ids * codec_flag) * codec_flag.unsqueeze(-1) face_embs = self.face_linear(face_embs) inputs_embeds = text_embeds + timespk_embeds + codec_embs inputs_embeds = torch.cat([ inputs_embeds[:, 0:1, :], # sos token face_embs, # face embeddings inputs_embeds[:, 1:, :] # inputs_embeds after sos ], dim=1) prompt_codec = batch.get("prompt_codec", None) if prompt_codec is not None: codec_emb = self.codec_embed(prompt_codec) inputs_embeds = torch.cat((inputs_embeds, codec_emb), dim=1) return inputs_embeds @torch.no_grad() def inference( self, data_in, data_lengths=None, key: list = None, **kwargs, ): uttid = key[0] inputs_emb = self.inference_prepare(data_in, **kwargs) logging.info(f"{uttid}: min length: {kwargs['min_length']}, max length: {kwargs['max_length']}") dtype = dtype_map[kwargs.get("llm_dtype", "fp32")] if not hasattr(self, "llm_generator"): llm_generator_conf = kwargs.get("dataset_conf", {}) self.llm_generator = LLMDecoder( token_embeder=self.codec_embed, **llm_generator_conf ).to(dtype) if (kwargs.get("use_qlora",False) or kwargs.get("infer_use_lora",False)) and (not kwargs.get("infer_lora_merged",False)): self.llm.base_model.model.lm_head = self.codec_head.to(dtype) else: self.llm.lm_head = self.codec_head.to(dtype) gen_codec, hit_eos, states = self.llm_generator( inputs_emb.to(dtype), self.llm, states=kwargs.get("states", {}), **kwargs ) output_dir = kwargs.get("output_dir", None) if output_dir is not None: output_dir = os.path.join(output_dir, "codec") os.makedirs(output_dir, exist_ok=True) np.save( os.path.join(output_dir, f"{key[0]}.npy"), gen_codec[0].cpu().numpy() ) return gen_codec, hit_eos, states