wehub-resource-sync
1005 lines
41 KiB
Python
1005 lines
41 KiB
Python
import logging
|
||
import os
|
||
import random
|
||
import re
|
||
import string
|
||
import time
|
||
import traceback
|
||
from typing import Union
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
|
||
from funasr.metrics.compute_acc import compute_accuracy
|
||
from funasr.register import tables
|
||
from funasr.train_utils.device_funcs import force_gatherable, to_device
|
||
from funasr.utils.datadir_writer import DatadirWriter
|
||
from funasr.utils.load_utils import extract_fbank, load_audio_text_image_video
|
||
try:
|
||
from transformers import AutoConfig, AutoModelForCausalLM
|
||
except ImportError:
|
||
AutoConfig = None
|
||
AutoModelForCausalLM = None
|
||
|
||
from .ctc import CTC
|
||
from .device_utils import resolve_autocast_device_type
|
||
from .tools.utils import forced_align
|
||
|
||
dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}
|
||
|
||
|
||
@tables.register("model_classes", "FunASRNano")
|
||
class FunASRNano(nn.Module):
|
||
"""Fun-ASR-Nano: End-to-End ASR Large Model.
|
||
|
||
Trained on tens of millions of hours of real speech data.
|
||
Supports 31 languages including Chinese dialects and regional accents.
|
||
|
||
Features:
|
||
- Character-level timestamps (via CTC forced alignment)
|
||
- Hotword customization
|
||
- Speaker diarization (when combined with spk_model)
|
||
- Lyrics and rap recognition
|
||
- Streaming chunk-by-chunk inference (demo2.py)
|
||
|
||
Output: {"key": ..., "text": ..., "timestamps": [{"token", "start_time", "end_time"}, ...],
|
||
"ctc_timestamps": [...]}
|
||
|
||
Note: Outputs punctuation natively — punc_model is NOT needed.
|
||
|
||
Requirements: pip install tiktoken huggingface_hub
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
audio_encoder: str = None,
|
||
audio_encoder_conf: dict = None,
|
||
audio_adaptor: str = None,
|
||
audio_adaptor_conf: dict = None,
|
||
llm: str = None,
|
||
llm_conf: dict = None,
|
||
input_size: int = 80,
|
||
length_normalized_loss: bool = False,
|
||
**kwargs,
|
||
):
|
||
"""Initialize FunASRNano.
|
||
|
||
Args:
|
||
audio_encoder: TODO.
|
||
audio_encoder_conf: Configuration dict for audio_encoder.
|
||
audio_adaptor: TODO.
|
||
audio_adaptor_conf: Configuration dict for audio_adaptor.
|
||
llm: TODO.
|
||
llm_conf: Configuration dict for llm.
|
||
input_size: Size/dimension parameter.
|
||
length_normalized_loss: TODO.
|
||
**kwargs: Additional keyword arguments.
|
||
"""
|
||
super().__init__()
|
||
|
||
# audio encoder
|
||
hub = audio_encoder_conf.get("hub", None)
|
||
self.audio_encoder_activation_checkpoint = audio_encoder_conf.get(
|
||
"activation_checkpoint", False
|
||
)
|
||
if hub == "ms":
|
||
from funasr import AutoModel
|
||
|
||
model = AutoModel(model=audio_encoder, model_revision="master")
|
||
audio_encoder_output_size = (
|
||
model.model.encoder_output_size
|
||
if hasattr(model.model, "encoder_output_size")
|
||
else -1
|
||
)
|
||
audio_encoder = (
|
||
model.model.model.encoder if hasattr(model.model, "model") else model.model.encoder
|
||
)
|
||
else:
|
||
encoder_class = tables.encoder_classes.get(audio_encoder)
|
||
audio_encoder = encoder_class(input_size=input_size, **audio_encoder_conf)
|
||
audio_encoder_output_size = audio_encoder.output_size()
|
||
freeze = audio_encoder_conf.get("freeze", True)
|
||
|
||
if freeze:
|
||
for _, param in audio_encoder.named_parameters():
|
||
param.requires_grad = False
|
||
audio_encoder.eval()
|
||
self.audio_encoder = audio_encoder
|
||
|
||
# llm
|
||
self.llm = None
|
||
init_param_path = llm_conf.get("init_param_path", None)
|
||
llm_dim = None
|
||
|
||
llm_load_kwargs = llm_conf.get("load_kwargs", {})
|
||
config = AutoConfig.from_pretrained(init_param_path)
|
||
model = AutoModelForCausalLM.from_config(config, **llm_load_kwargs)
|
||
|
||
freeze = llm_conf.get("freeze", True)
|
||
if freeze:
|
||
for _, param in model.named_parameters():
|
||
param.requires_grad = False
|
||
model.eval()
|
||
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]
|
||
|
||
# adaptor
|
||
adaptor_class = tables.adaptor_classes.get(audio_adaptor)
|
||
if audio_encoder_output_size > 0:
|
||
audio_adaptor_conf["encoder_dim"] = audio_encoder_output_size
|
||
audio_adaptor_conf["llm_dim"] = (
|
||
llm_dim if llm_dim is not None else audio_adaptor_conf["llm_dim"]
|
||
)
|
||
audio_adaptor = adaptor_class(**audio_adaptor_conf)
|
||
freeze = audio_adaptor_conf.get("freeze", False)
|
||
if freeze:
|
||
for _, param in audio_adaptor.named_parameters():
|
||
param.requires_grad = False
|
||
audio_adaptor.eval()
|
||
self.audio_adaptor = audio_adaptor
|
||
self.use_low_frame_rate = audio_adaptor_conf.get("use_low_frame_rate", False)
|
||
|
||
# ctc decoder
|
||
self.ctc_decoder = None
|
||
# TODO: fix table name
|
||
ctc_decoder_class = tables.adaptor_classes.get(kwargs.get("ctc_decoder", None))
|
||
if ctc_decoder_class is not None:
|
||
ctc_tokenizer = (
|
||
kwargs.get("ctc_tokenizer", None)
|
||
if "ctc_tokenizer" in kwargs
|
||
else kwargs["dataset_conf"]["ctc_tokenizer"]
|
||
)
|
||
ctc_tokenizer_conf = (
|
||
kwargs.get("ctc_tokenizer_conf", None)
|
||
if "ctc_tokenizer_conf" in kwargs
|
||
else kwargs["dataset_conf"]["ctc_tokenizer_conf"]
|
||
)
|
||
if ctc_tokenizer is not None and ctc_tokenizer_conf is not None:
|
||
ctc_tokenizer_class = tables.tokenizer_classes.get(ctc_tokenizer)
|
||
ctc_tokenizer = ctc_tokenizer_class(**ctc_tokenizer_conf)
|
||
self.ctc_tokenizer = ctc_tokenizer
|
||
assert ctc_tokenizer is not None, f"ctc_tokenizer must be set"
|
||
|
||
ctc_vocab_size = kwargs.get("ctc_vocab_size", 60515)
|
||
ctc_decoder_conf = kwargs.get("ctc_decoder_conf", {})
|
||
if audio_encoder_output_size > 0:
|
||
ctc_decoder_conf["encoder_dim"] = audio_encoder_output_size
|
||
self.ctc_decoder = ctc_decoder_class(**ctc_decoder_conf)
|
||
init_param_path = ctc_decoder_conf.get("init_param_path", None)
|
||
if init_param_path is not None:
|
||
src_state = torch.load(init_param_path, map_location="cpu")
|
||
flag = self.ctc_decoder.load_state_dict(src_state, strict=False)
|
||
logging.info(f"Loading ctc_decoder ckpt: {init_param_path}, status: {flag}")
|
||
freeze = ctc_decoder_conf.get("freeze", False)
|
||
if freeze:
|
||
for _, param in self.ctc_decoder.named_parameters():
|
||
param.requires_grad = False
|
||
self.ctc_decoder.eval()
|
||
|
||
ctc_conf = kwargs.get("ctc_conf", {})
|
||
self.blank_id = ctc_conf.get("blank_id", ctc_vocab_size - 1)
|
||
self.ctc_weight = kwargs.get("ctc_weight", 0.3)
|
||
self.ctc = CTC(
|
||
odim=ctc_vocab_size,
|
||
encoder_output_size=audio_encoder_output_size,
|
||
blank_id=self.blank_id,
|
||
**ctc_conf,
|
||
)
|
||
self.detach_ctc_decoder = kwargs.get("detach_ctc_decoder", True)
|
||
self.error_calculator = None
|
||
|
||
self.length_normalized_loss = length_normalized_loss
|
||
rank = int(os.environ.get("RANK", 0))
|
||
logging.info(f"rank: {rank}, model is builded.")
|
||
|
||
def forward(
|
||
self,
|
||
speech: torch.Tensor = None,
|
||
speech_lengths: torch.Tensor = None,
|
||
input_ids: torch.Tensor = None,
|
||
attention_mask: torch.Tensor = None,
|
||
labels_ids: torch.Tensor = None,
|
||
fbank_beg: torch.Tensor = None,
|
||
fbank_mask: torch.Tensor = None,
|
||
**kwargs,
|
||
):
|
||
"""Forward pass for training.
|
||
|
||
Args:
|
||
speech: Speech audio tensor, shape (batch, time).
|
||
speech_lengths: Length of each speech sample.
|
||
input_ids: TODO.
|
||
attention_mask: TODO.
|
||
labels_ids: TODO.
|
||
fbank_beg: TODO.
|
||
fbank_mask: TODO.
|
||
**kwargs: Additional keyword arguments.
|
||
"""
|
||
batch_size, token_num = input_ids.shape
|
||
stats = {}
|
||
input_ids[input_ids < 0] = 0
|
||
inputs_embeds = self.llm.model.get_input_embeddings()(input_ids)
|
||
if speech is not None:
|
||
if len(speech_lengths.size()) > 1:
|
||
speech_lengths = speech_lengths[:, 0]
|
||
batch_size_speech, frames, _ = speech.shape
|
||
|
||
# audio encoder
|
||
if self.audio_encoder_activation_checkpoint:
|
||
from torch.utils.checkpoint import checkpoint
|
||
|
||
encoder_out, encoder_out_lens = checkpoint(
|
||
self.encode, speech, speech_lengths, use_reentrant=False
|
||
)
|
||
else:
|
||
encoder_out, encoder_out_lens = self.encode(speech, speech_lengths)
|
||
|
||
# audio_adaptor
|
||
encoder_out, encoder_out_lens = self.audio_adaptor(encoder_out, encoder_out_lens)
|
||
|
||
batch_size, token_num, dims = inputs_embeds.shape
|
||
fake_token_len = kwargs.get("fake_token_len")
|
||
fake_token_len[fake_token_len < 0] = 0
|
||
fbank_beg[fbank_beg < 0] = 0
|
||
|
||
speech_idx = 0
|
||
for batch_idx in range(batch_size):
|
||
for turn_id in range(fbank_beg.shape[1]):
|
||
fbank_beg_idx = fbank_beg[batch_idx, turn_id].item()
|
||
if fbank_beg_idx > 0:
|
||
speech_token_len = fake_token_len[batch_idx, turn_id]
|
||
speech_token = encoder_out[speech_idx, :speech_token_len, :]
|
||
|
||
try:
|
||
inputs_embeds[
|
||
batch_idx,
|
||
fbank_beg_idx : fbank_beg_idx + speech_token_len,
|
||
:,
|
||
] = speech_token
|
||
except Exception as e:
|
||
logging.error(f"{str(e)}, {traceback.format_exc()}")
|
||
logging.info(
|
||
f"batch_idx: {batch_idx}, inputs_embeds: {inputs_embeds.shape}, fbank_beg_idx: {fbank_beg_idx}, speech_token_len: {speech_token_len}, encoder_out: {encoder_out.shape}, encoder_out_lens: {encoder_out_lens}, fake_token_len: {fake_token_len}, speech_lengths: {speech_lengths}"
|
||
)
|
||
speech_token_len = encoder_out_lens[speech_idx].item()
|
||
speech_token = encoder_out[speech_idx, :speech_token_len, :]
|
||
inputs_embeds[
|
||
batch_idx,
|
||
fbank_beg_idx : fbank_beg_idx + speech_token_len,
|
||
:,
|
||
] = speech_token
|
||
|
||
speech_idx += 1
|
||
|
||
stats["batch_size_speech"] = batch_size_speech
|
||
stats["batch_size_x_frames"] = frames * batch_size_speech
|
||
stats["batch_size_real_frames"] = speech_lengths.sum().item()
|
||
stats["padding_frames"] = stats["batch_size_x_frames"] - stats["batch_size_real_frames"]
|
||
|
||
autocast_device_type = resolve_autocast_device_type(next(self.parameters()).device)
|
||
with torch.autocast(
|
||
device_type=autocast_device_type,
|
||
enabled=True if self.llm_dtype != "fp32" else False,
|
||
dtype=dtype_map[self.llm_dtype],
|
||
):
|
||
labels_ids[labels_ids == -1] = -100
|
||
attention_mask[attention_mask < 0] = 0
|
||
model_outputs = self.llm(
|
||
inputs_embeds=inputs_embeds.to(dtype_map[self.llm_dtype]),
|
||
attention_mask=attention_mask,
|
||
labels=labels_ids,
|
||
)
|
||
loss = model_outputs.loss
|
||
|
||
with torch.no_grad():
|
||
preds = torch.argmax(model_outputs.logits, -1)
|
||
acc_att = compute_accuracy(preds[:, :-1], labels_ids[:, 1:], ignore_label=-100)
|
||
stats["acc"] = acc_att
|
||
|
||
stats["loss"] = torch.clone(loss.detach())
|
||
stats["batch_size"] = batch_size
|
||
|
||
stats["batch_size_x_tokens"] = token_num * batch_size
|
||
stats["batch_size_real_tokens"] = attention_mask.sum().item()
|
||
stats["padding_tokens"] = stats["batch_size_x_tokens"] - stats["batch_size_real_tokens"]
|
||
|
||
dialog_turns = (fbank_beg > 0).sum(-1)
|
||
dialog_turns_max = torch.max(dialog_turns).int().item()
|
||
dialog_turns_avg = dialog_turns.sum().item() / batch_size
|
||
stats["dialog_turns_max"] = dialog_turns_max
|
||
stats["dialog_turns_avg"] = dialog_turns_avg
|
||
|
||
# force_gatherable: to-device and to-tensor if scalar for DataParallel
|
||
if self.length_normalized_loss:
|
||
batch_size = int((labels_ids > 0 + 1).sum())
|
||
loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device)
|
||
return loss, stats, weight
|
||
|
||
def forward_export(self, speech, speech_lengths, **kwargs):
|
||
"""Forward export.
|
||
|
||
Args:
|
||
speech: Speech audio tensor, shape (batch, time).
|
||
speech_lengths: Length of each speech sample.
|
||
**kwargs: Additional keyword arguments.
|
||
"""
|
||
x, olens = self.audio_encoder(speech, speech_lengths)
|
||
encoder_out, encoder_out_lens = self.audio_adaptor(x, olens)
|
||
return encoder_out, encoder_out_lens
|
||
|
||
def encode(self, speech, speech_lengths):
|
||
# audio encoder
|
||
"""Encode.
|
||
|
||
Args:
|
||
speech: Speech audio tensor, shape (batch, time).
|
||
speech_lengths: Length of each speech sample.
|
||
"""
|
||
encoder_out, encoder_out_lens = self.audio_encoder(speech, speech_lengths)
|
||
|
||
return encoder_out, encoder_out_lens
|
||
|
||
def data_template(self, data):
|
||
"""Data template.
|
||
|
||
Args:
|
||
data: TODO.
|
||
"""
|
||
system, user, assistant = [], [], []
|
||
for i, item in enumerate(data):
|
||
role = item["role"]
|
||
content = item["content"]
|
||
if role == "system":
|
||
system.append(content)
|
||
elif role == "user":
|
||
if "audio" in item:
|
||
audio = item["audio"]
|
||
content = [content, audio]
|
||
user.append(content)
|
||
elif role == "assistant":
|
||
assistant.append(content)
|
||
|
||
system = system * len(user)
|
||
|
||
contents = {
|
||
"system": system,
|
||
"user": user,
|
||
"assistant": assistant,
|
||
}
|
||
|
||
return contents
|
||
|
||
def data_load_speech(self, contents: dict, tokenizer, frontend, meta_data={}, **kwargs):
|
||
"""Data load speech.
|
||
|
||
Args:
|
||
contents: TODO.
|
||
tokenizer: Tokenizer instance for text encoding/decoding.
|
||
frontend: Audio frontend for feature extraction.
|
||
meta_data: TODO.
|
||
**kwargs: Additional keyword arguments.
|
||
"""
|
||
system = contents["system"]
|
||
user = contents["user"]
|
||
assistant = contents["assistant"]
|
||
pattern = re.compile(r"(<\|startofspeech\|>.*?<\|endofspeech\|>)")
|
||
do_think = True
|
||
sys_prompt = True
|
||
if "dataset_conf" in kwargs:
|
||
do_think = kwargs["dataset_conf"].get("do_think", True)
|
||
sys_prompt = kwargs["dataset_conf"].get("sys_prompt", True)
|
||
|
||
input_ids, labels, fbank, fbank_lens, fbank_mask, fbank_beg, fake_token_len = (
|
||
[],
|
||
[],
|
||
[],
|
||
[],
|
||
[],
|
||
[],
|
||
[],
|
||
)
|
||
input_source_ids = []
|
||
for i, (system_prompt, user_prompt, target_out) in enumerate(zip(system, user, assistant)):
|
||
if i >= kwargs.get("multiturn_num_max", 5):
|
||
break
|
||
if len(input_ids) > kwargs.get("max_token_length", 1500):
|
||
break
|
||
if isinstance(user_prompt, (list, tuple)):
|
||
user_prompt, audio = user_prompt
|
||
if i == 0:
|
||
if kwargs.get("infer_with_assistant_input", False):
|
||
source_input = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{user_prompt}"
|
||
if not sys_prompt:
|
||
source_input = f"<|im_start|>user\n{user_prompt}"
|
||
else:
|
||
source_input = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
|
||
if not sys_prompt:
|
||
source_input = (
|
||
f"<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
|
||
)
|
||
else:
|
||
if kwargs.get("infer_with_assistant_input", False):
|
||
source_input = f"<|im_start|>user\n{user_prompt}"
|
||
else:
|
||
source_input = (
|
||
f"<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
|
||
)
|
||
if not do_think:
|
||
source_input += "<think>\n\n</think>\n\n"
|
||
if kwargs.get("prev_text", None) is not None:
|
||
source_input += kwargs["prev_text"]
|
||
|
||
splits = pattern.split(source_input)
|
||
source_ids = []
|
||
fbank_mask_i = []
|
||
fake_token_len_i = 0
|
||
fbank_beg_i = -1
|
||
speech, speech_lengths = [], []
|
||
for k, sub_str in enumerate(splits):
|
||
if not sub_str.startswith("<|startofspeech|>"):
|
||
sub_token = tokenizer.encode(sub_str)
|
||
source_ids += sub_token
|
||
fbank_mask_i += [0] * len(sub_token)
|
||
else:
|
||
sub_str = sub_str.replace("<|startofspeech|>", "").replace(
|
||
"<|endofspeech|>", ""
|
||
)
|
||
if sub_str.startswith("!"):
|
||
sub_str = sub_str[1:]
|
||
if sub_str.startswith("!"): # !!: audio sample point
|
||
sub_str = audio
|
||
try:
|
||
time1 = time.perf_counter()
|
||
data_src = load_audio_text_image_video(
|
||
sub_str, fs=frontend.fs, **kwargs
|
||
)
|
||
time2 = time.perf_counter()
|
||
meta_data["load_data"] = f"{time2 - time1:0.3f}"
|
||
except Exception as e:
|
||
logging.error(f"Loading wav failed! {str(e)}, {traceback.format_exc()}")
|
||
|
||
speech, speech_lengths = extract_fbank(
|
||
data_src,
|
||
data_type=kwargs.get("data_type", "sound"),
|
||
frontend=frontend,
|
||
is_final=True,
|
||
) # speech: [b, T, d]
|
||
|
||
time3 = time.perf_counter()
|
||
meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
|
||
meta_data["batch_data_time"] = (
|
||
speech_lengths.sum().item()
|
||
* frontend.frame_shift
|
||
* frontend.lfr_n
|
||
/ 1000
|
||
)
|
||
|
||
if self.use_low_frame_rate:
|
||
olens = 1 + (speech_lengths[0].item() - 3 + 2 * 1) // 2
|
||
olens = 1 + (olens - 3 + 2 * 1) // 2
|
||
fake_token_len_i = (olens - 1) // 2 + 1
|
||
else:
|
||
fake_token_len_i = speech_lengths[0].item()
|
||
fake_token = [0] * fake_token_len_i
|
||
fbank_beg_i = len(source_ids)
|
||
source_ids += fake_token
|
||
fbank_mask_i += [1] * len(fake_token)
|
||
|
||
fbank_beg += [fbank_beg_i + len(input_ids)]
|
||
fake_token_len += [fake_token_len_i]
|
||
source_mask = [-100] * len(source_ids)
|
||
target_out = f"{target_out}<|im_end|>"
|
||
target_ids = tokenizer.encode(target_out)
|
||
input_source_ids = input_ids + source_ids
|
||
input_ids += source_ids + target_ids
|
||
labels += source_mask + target_ids
|
||
fbank_mask += fbank_mask_i
|
||
if len(speech) > 0:
|
||
fbank.append(speech[0, :, :])
|
||
fbank_lens.append(speech_lengths)
|
||
|
||
input_ids = torch.tensor(input_ids, dtype=torch.int64) # [: self.max_token_length]
|
||
attention_mask = torch.tensor([1] * len(input_ids), dtype=torch.int32)
|
||
labels = torch.tensor(labels, dtype=torch.int64) # [: self.max_token_length]
|
||
|
||
fbank_mask = torch.tensor(fbank_mask, dtype=torch.float32)
|
||
fbank_beg = torch.tensor(fbank_beg, dtype=torch.int32)
|
||
fake_token_len = torch.tensor(fake_token_len, dtype=torch.int32)
|
||
source_ids = torch.tensor(input_source_ids, dtype=torch.int64)
|
||
target_ids = torch.tensor(target_ids, dtype=torch.int64)
|
||
|
||
if len(fbank) > 0:
|
||
speech = torch.nn.utils.rnn.pad_sequence(fbank, batch_first=True, padding_value=0.0)
|
||
speech_lengths = torch.nn.utils.rnn.pad_sequence(
|
||
fbank_lens, batch_first=True, padding_value=-1
|
||
)
|
||
else:
|
||
speech = []
|
||
speech_lengths = []
|
||
output = {
|
||
"speech": speech,
|
||
"speech_lengths": speech_lengths,
|
||
"fbank_mask": fbank_mask[None, :],
|
||
"fbank_beg": fbank_beg[None,],
|
||
"fake_token_len": fake_token_len[None, :],
|
||
"input_ids": input_ids[None,],
|
||
"attention_mask": attention_mask[None,],
|
||
"labels_ids": labels,
|
||
"source_ids": source_ids[None, :],
|
||
"target_ids": target_ids[None, :],
|
||
}
|
||
|
||
return output
|
||
|
||
def inference_prepare(
|
||
self,
|
||
data_in,
|
||
data_lengths=None,
|
||
key: list = None,
|
||
tokenizer=None,
|
||
frontend=None,
|
||
**kwargs,
|
||
):
|
||
"""Inference prepare.
|
||
|
||
Args:
|
||
data_in: Input data (audio samples, file paths, or text).
|
||
data_lengths: Lengths of each input sample in the batch.
|
||
key: Sample identifiers.
|
||
tokenizer: Tokenizer instance for text encoding/decoding.
|
||
frontend: Audio frontend for feature extraction.
|
||
**kwargs: Additional keyword arguments.
|
||
"""
|
||
meta_data = {}
|
||
|
||
if len(data_in) > 1:
|
||
raise NotImplementedError("batch decoding is not implemented")
|
||
|
||
contents = self.data_template(data_in[0])
|
||
output = self.data_load_speech(contents, tokenizer, frontend, meta_data=meta_data, **kwargs)
|
||
batch = to_device(output, kwargs["device"])
|
||
|
||
# audio encoder
|
||
speech = batch["speech"]
|
||
|
||
if len(speech) > 0:
|
||
if "audio_embedding" in kwargs and "audio_embedding_lens" in kwargs:
|
||
encoder_out = kwargs["audio_embedding"]
|
||
encoder_out_lens = kwargs["audio_embedding_lens"]
|
||
else:
|
||
speech_lengths = batch["speech_lengths"][:, 0]
|
||
# NOTE: the audio encoder contains fp32-only ops, so casting its
|
||
# input to fp16/bf16 here raises a dtype mismatch. The encoder
|
||
# therefore always runs in fp32; low precision (fp16/bf16) is
|
||
# applied to the LLM decoder only. The audio embeddings are cast
|
||
# to the LLM dtype automatically when written into inputs_embeds.
|
||
# audio encoder
|
||
encoder_out, encoder_out_lens = self.encode(speech, speech_lengths)
|
||
|
||
# audio_adaptor
|
||
adaptor_out, adaptor_out_lens = self.audio_adaptor(encoder_out, encoder_out_lens)
|
||
meta_data["encoder_out"] = encoder_out
|
||
meta_data["encoder_out_lens"] = encoder_out_lens
|
||
meta_data["audio_adaptor_out"] = adaptor_out
|
||
meta_data["audio_adaptor_out_lens"] = adaptor_out_lens
|
||
|
||
input_ids = batch["input_ids"]
|
||
source_ids = batch["source_ids"]
|
||
fbank_beg = batch["fbank_beg"]
|
||
fake_token_len = batch["fake_token_len"]
|
||
|
||
if not kwargs.get("teacherforcing", False):
|
||
input_ids = source_ids
|
||
|
||
input_ids[input_ids < 0] = 0
|
||
inputs_embeds = self.llm.model.get_input_embeddings()(input_ids)
|
||
|
||
batch_size, token_num, dims = inputs_embeds.shape
|
||
|
||
fake_token_len[fake_token_len < 0] = 0
|
||
fbank_beg[fbank_beg < 0] = 0
|
||
|
||
speech_idx = 0
|
||
for batch_idx in range(batch_size):
|
||
for turn_id in range(fbank_beg.shape[1]):
|
||
fbank_beg_idx = fbank_beg[batch_idx, turn_id].item()
|
||
if fbank_beg_idx > 0:
|
||
speech_token_len = fake_token_len[batch_idx, turn_id]
|
||
speech_token = adaptor_out[speech_idx, :speech_token_len, :]
|
||
|
||
try:
|
||
inputs_embeds[
|
||
batch_idx,
|
||
fbank_beg_idx : fbank_beg_idx + speech_token_len,
|
||
:,
|
||
] = speech_token
|
||
except Exception as e:
|
||
#
|
||
logging.error(f"{str(e)}, {traceback.format_exc()}")
|
||
logging.info(
|
||
f"batch_idx: {batch_idx}, inputs_embeds: {inputs_embeds.shape}, fbank_beg_idx: {fbank_beg_idx}, speech_token_len: {speech_token_len}, adaptor_out: {adaptor_out.shape}, adaptor_out_lens: {adaptor_out_lens}, fake_token_len: {fake_token_len}, speech_lengths: {speech_lengths}"
|
||
)
|
||
speech_token_len = adaptor_out_lens[speech_idx].item()
|
||
speech_token = adaptor_out[speech_idx, :speech_token_len, :]
|
||
inputs_embeds[
|
||
batch_idx,
|
||
fbank_beg_idx : fbank_beg_idx + speech_token_len,
|
||
:,
|
||
] = speech_token
|
||
|
||
speech_idx += 1
|
||
return inputs_embeds, contents, batch, source_ids, meta_data
|
||
|
||
def get_prompt(self, hotwords: list[str], language: str = None, itn: bool = True):
|
||
"""Get prompt.
|
||
|
||
Args:
|
||
hotwords: TODO.
|
||
language: Language identifier.
|
||
itn: TODO.
|
||
"""
|
||
if len(hotwords) > 0:
|
||
hotwords = ", ".join(hotwords)
|
||
prompt = f"请结合上下文信息，更加准确地完成语音转写任务。如果没有相关信息，我们会留空。\n\n\n**上下文信息：**\n\n\n"
|
||
prompt += f"热词列表：[{hotwords}]\n"
|
||
else:
|
||
prompt = ""
|
||
if language is None:
|
||
prompt += "语音转写"
|
||
else:
|
||
prompt += f"语音转写成{language}"
|
||
if not itn:
|
||
prompt += "，不进行文本规整"
|
||
return prompt + "："
|
||
|
||
def generate_chatml(self, prompt: str, data: Union[str, torch.Tensor]):
|
||
"""Generate chatml.
|
||
|
||
Args:
|
||
prompt: TODO.
|
||
data: TODO.
|
||
"""
|
||
if isinstance(data, str):
|
||
return [
|
||
{"role": "system", "content": "You are a helpful assistant."},
|
||
{"role": "user", "content": f"{prompt}<|startofspeech|>!{data}<|endofspeech|>"},
|
||
{"role": "assistant", "content": "null"},
|
||
]
|
||
elif isinstance(data, torch.Tensor):
|
||
return [
|
||
{"role": "system", "content": "You are a helpful assistant."},
|
||
{
|
||
"role": "user",
|
||
"content": f"{prompt}<|startofspeech|>!!<|endofspeech|>",
|
||
"audio": data,
|
||
},
|
||
{"role": "assistant", "content": "null"},
|
||
]
|
||
|
||
def inference(
|
||
self,
|
||
data_in,
|
||
data_lengths=None,
|
||
key: list = None,
|
||
tokenizer=None,
|
||
frontend=None,
|
||
**kwargs,
|
||
):
|
||
"""Run inference on input data.
|
||
|
||
Args:
|
||
data_in: Input data (audio samples, file paths, or text).
|
||
data_lengths: Lengths of each input sample in the batch.
|
||
key: Sample identifiers.
|
||
tokenizer: Tokenizer instance for text encoding/decoding.
|
||
frontend: Audio frontend for feature extraction.
|
||
**kwargs: Additional keyword arguments.
|
||
"""
|
||
prompt = self.get_prompt(
|
||
kwargs.get("hotwords", []), kwargs.get("language", None), kwargs.get("itn", True)
|
||
)
|
||
data_in = [self.generate_chatml(prompt, data) for data in data_in]
|
||
|
||
if key is None:
|
||
key = []
|
||
for _ in data_in:
|
||
chars = string.ascii_letters + string.digits
|
||
key.append("rand_key_" + "".join(random.choice(chars) for _ in range(13)))
|
||
|
||
return self.inference_llm(
|
||
data_in,
|
||
data_lengths=data_lengths,
|
||
key=key,
|
||
tokenizer=tokenizer,
|
||
frontend=frontend,
|
||
**kwargs,
|
||
)
|
||
|
||
def _inference_llm_batch(self, data_in, data_lengths, key, tokenizer, frontend, **kwargs):
|
||
"""Batched LLM decoding for multiple VAD segments at once.
|
||
|
||
Builds each segment's inputs_embeds via the single-sample
|
||
inference_prepare, left-pads them into one batch, and runs a single
|
||
llm.generate. This greatly improves GPU utilization for the small LLM
|
||
decoder (the per-segment, batch_size=1 path underuses the GPU).
|
||
CTC timestamps are not produced in batched mode.
|
||
"""
|
||
# normalize nested key (e.g. [[k1, k2, ...]]) like the single-sample path
|
||
if key is not None and len(key) > 0 and isinstance(key[0], (list, tuple)):
|
||
key = list(key[0])
|
||
embs = []
|
||
keys = []
|
||
for i, d in enumerate(data_in):
|
||
k_i = [key[i]] if key is not None and i < len(key) else None
|
||
emb_i, _c, _b, _s, _m = self.inference_prepare(
|
||
[d], data_lengths, k_i, tokenizer, frontend, **kwargs
|
||
)
|
||
embs.append(emb_i)
|
||
keys.append(key[i] if key is not None and i < len(key) else f"rand_{i}")
|
||
|
||
llm_dtype = kwargs.get("llm_dtype", "fp32")
|
||
if llm_dtype == "fp32":
|
||
llm_dtype = "fp16" if kwargs.get("fp16", False) else llm_dtype
|
||
llm_dtype = "bf16" if kwargs.get("bf16", False) else llm_dtype
|
||
dt = dtype_map[llm_dtype]
|
||
device = embs[0].device
|
||
self.llm = self.llm.to(dt)
|
||
|
||
B = len(embs)
|
||
D = embs[0].shape[-1]
|
||
Tmax = max(e.shape[1] for e in embs)
|
||
padded = torch.zeros(B, Tmax, D, device=device, dtype=dt)
|
||
attn = torch.zeros(B, Tmax, dtype=torch.long, device=device)
|
||
for i, e in enumerate(embs):
|
||
Ti = e.shape[1]
|
||
padded[i, Tmax - Ti :, :] = e[0].to(dt) # left padding
|
||
attn[i, Tmax - Ti :] = 1
|
||
|
||
autocast_device_type = resolve_autocast_device_type(kwargs.get("device", "cuda"))
|
||
with torch.autocast(
|
||
device_type=autocast_device_type,
|
||
enabled=True if llm_dtype != "fp32" else False,
|
||
dtype=dt,
|
||
):
|
||
# left padding requires explicit position_ids so each segment's real
|
||
# tokens get positions 0,1,2,... regardless of the padding length.
|
||
position_ids = attn.long().cumsum(-1) - 1
|
||
position_ids.masked_fill_(attn == 0, 1)
|
||
generated_ids = self.llm.generate(
|
||
inputs_embeds=padded,
|
||
attention_mask=attn,
|
||
position_ids=position_ids,
|
||
max_new_tokens=kwargs.get("max_length", 512),
|
||
pad_token_id=(
|
||
self.llm.config.pad_token_id
|
||
if self.llm.config.pad_token_id is not None
|
||
else self.llm.config.eos_token_id
|
||
),
|
||
**kwargs.get("llm_kwargs", {}),
|
||
)
|
||
texts = tokenizer.batch_decode(
|
||
generated_ids, skip_special_tokens=kwargs.get("skip_special_tokens", True)
|
||
)
|
||
results = []
|
||
for i, t in enumerate(texts):
|
||
t = kwargs.get("prev_text", "") + t
|
||
results.append(
|
||
{
|
||
"key": keys[i],
|
||
"text": re.sub(r"\s+", " ", t.replace("/sil", " ")),
|
||
"text_tn": re.sub(r"[^\w\s\u3000\u4e00-\u9fff]+", "", t),
|
||
}
|
||
)
|
||
return results, {}
|
||
|
||
@staticmethod
|
||
def _slice_batch_value(value, index):
|
||
if value is None:
|
||
return None
|
||
if isinstance(value, torch.Tensor) and value.ndim > 0 and value.shape[0] > index:
|
||
return value[index : index + 1]
|
||
if isinstance(value, list) and len(value) > index:
|
||
return [value[index]]
|
||
if isinstance(value, tuple) and len(value) > index:
|
||
return (value[index],)
|
||
return value
|
||
|
||
@staticmethod
|
||
def _merge_inference_meta(target, source):
|
||
for name, value in source.items():
|
||
if isinstance(value, (int, float)):
|
||
target[name] = target.get(name, 0.0) + value
|
||
elif name not in target:
|
||
target[name] = value
|
||
|
||
def _inference_llm_ctc_sequential(
|
||
self, data_in, data_lengths, key, tokenizer, frontend, **kwargs
|
||
):
|
||
"""Run multi-segment input one segment at a time when CTC timestamps are active."""
|
||
if key is not None and len(key) > 0 and isinstance(key[0], (list, tuple)):
|
||
key = list(key[0])
|
||
|
||
results = []
|
||
meta_data = {}
|
||
for i, data_i in enumerate(data_in):
|
||
key_i = [key[i]] if key is not None and i < len(key) else None
|
||
data_lengths_i = self._slice_batch_value(data_lengths, i)
|
||
results_i, meta_i = self.inference_llm(
|
||
[data_i],
|
||
data_lengths=data_lengths_i,
|
||
key=key_i,
|
||
tokenizer=tokenizer,
|
||
frontend=frontend,
|
||
**kwargs,
|
||
)
|
||
results.extend(results_i)
|
||
self._merge_inference_meta(meta_data, meta_i)
|
||
return results, meta_data
|
||
|
||
def inference_llm(
|
||
self,
|
||
data_in,
|
||
data_lengths=None,
|
||
key: list = None,
|
||
tokenizer=None,
|
||
frontend=None,
|
||
**kwargs,
|
||
):
|
||
"""Inference llm.
|
||
|
||
Args:
|
||
data_in: Input data (audio samples, file paths, or text).
|
||
data_lengths: Lengths of each input sample in the batch.
|
||
key: Sample identifiers.
|
||
tokenizer: Tokenizer instance for text encoding/decoding.
|
||
frontend: Audio frontend for feature extraction.
|
||
**kwargs: Additional keyword arguments.
|
||
"""
|
||
# Only batch when CTC timestamps are not needed; the batched path does not
|
||
# produce ctc_timestamps, so fall back to the single-sample path when a CTC
|
||
# decoder is loaded (preserves timestamp behavior).
|
||
if len(data_in) > 1:
|
||
if self.ctc_decoder is None:
|
||
return self._inference_llm_batch(
|
||
data_in, data_lengths, key, tokenizer, frontend, **kwargs
|
||
)
|
||
return self._inference_llm_ctc_sequential(
|
||
data_in, data_lengths, key, tokenizer, frontend, **kwargs
|
||
)
|
||
inputs_embeds, contents, batch, source_ids, meta_data = self.inference_prepare(
|
||
data_in, data_lengths, key, tokenizer, frontend, **kwargs
|
||
)
|
||
|
||
ctc_results = []
|
||
if self.ctc_decoder is not None:
|
||
encoder_out = meta_data["encoder_out"]
|
||
encoder_out_lens = meta_data["encoder_out_lens"]
|
||
decoder_out, decoder_out_lens = self.ctc_decoder(encoder_out, encoder_out_lens)
|
||
ctc_logits = self.ctc.log_softmax(decoder_out)
|
||
|
||
b, n, d = encoder_out.size()
|
||
if isinstance(key[0], (list, tuple)):
|
||
key = key[0]
|
||
if len(key) < b:
|
||
key = key * b
|
||
for i in range(b):
|
||
x = ctc_logits[i, : encoder_out_lens[i].item(), :]
|
||
yseq = x.argmax(dim=-1)
|
||
yseq = torch.unique_consecutive(yseq, dim=-1)
|
||
mask = yseq != self.blank_id
|
||
token_int = yseq[mask].tolist()
|
||
# Change integer-ids to tokens
|
||
text = self.ctc_tokenizer.decode(token_int)
|
||
ctc_results.append({"key": key[i], "text": text, "ctc_logits": x})
|
||
|
||
llm_dtype = kwargs.get("llm_dtype", "fp32")
|
||
if llm_dtype == "fp32":
|
||
llm_dtype = "fp16" if kwargs.get("fp16", False) else llm_dtype
|
||
llm_dtype = "bf16" if kwargs.get("bf16", False) else llm_dtype
|
||
|
||
autocast_device_type = resolve_autocast_device_type(kwargs.get("device", "cuda"))
|
||
with torch.autocast(
|
||
device_type=autocast_device_type,
|
||
enabled=True if llm_dtype != "fp32" else False,
|
||
dtype=dtype_map[llm_dtype],
|
||
):
|
||
label = contents["assistant"][-1]
|
||
self.llm = self.llm.to(dtype_map[llm_dtype])
|
||
inputs_embeds = inputs_embeds.to(dtype_map[llm_dtype])
|
||
llm_kwargs = kwargs.get("llm_kwargs", {})
|
||
if not kwargs.get("teacherforcing", False):
|
||
attention_mask = batch.get("attention_mask", None)
|
||
generated_ids = self.llm.generate(
|
||
inputs_embeds=inputs_embeds,
|
||
attention_mask=attention_mask,
|
||
max_new_tokens=kwargs.get("max_length", 512),
|
||
pad_token_id=self.llm.config.pad_token_id or self.llm.config.eos_token_id,
|
||
**llm_kwargs,
|
||
)
|
||
|
||
response = tokenizer.batch_decode(
|
||
generated_ids,
|
||
skip_special_tokens=kwargs.get("skip_special_tokens", True),
|
||
)[0]
|
||
|
||
loss = None
|
||
else:
|
||
labels_ids = batch["labels_ids"]
|
||
labels_ids[labels_ids == -1] = -100
|
||
attention_mask = batch.get("attention_mask", None)
|
||
model_outputs = self.llm(
|
||
inputs_embeds=inputs_embeds,
|
||
attention_mask=attention_mask,
|
||
labels=labels_ids,
|
||
pad_token_id=self.llm.config.pad_token_id or self.llm.config.eos_token_id,
|
||
**llm_kwargs,
|
||
)
|
||
|
||
preds = torch.argmax(model_outputs.logits, -1)[:, source_ids.shape[1] :]
|
||
response = tokenizer.batch_decode(
|
||
preds,
|
||
add_special_tokens=False,
|
||
skip_special_tokens=kwargs.get("skip_special_tokens", True),
|
||
)[0]
|
||
loss = model_outputs.loss.item()
|
||
response = kwargs.get("prev_text", "") + response
|
||
|
||
ibest_writer = None
|
||
if kwargs.get("output_dir") is not None:
|
||
if not hasattr(self, "writer"):
|
||
self.writer = DatadirWriter(kwargs.get("output_dir"))
|
||
ibest_writer = self.writer[f"{0 + 1}best_recog"]
|
||
|
||
results = []
|
||
response_clean = re.sub(r"[^\w\s\u3000\u4e00-\u9fff]+", "", response)
|
||
result_i = {
|
||
"key": key[0],
|
||
"text": re.sub(r"\s+", " ", response.replace("/sil", " ")),
|
||
"text_tn": response_clean,
|
||
"label": label,
|
||
}
|
||
if loss is not None:
|
||
result_i["loss"] = loss
|
||
results.append(result_i)
|
||
|
||
for ctc_result, result in zip(ctc_results, results):
|
||
result["ctc_text"] = ctc_result["text"].replace("<|nospeech|>", "")
|
||
target_ids = torch.tensor(
|
||
self.ctc_tokenizer.encode(result["ctc_text"]), dtype=torch.int64
|
||
)
|
||
result["ctc_timestamps"] = forced_align(
|
||
ctc_result["ctc_logits"], target_ids, self.blank_id
|
||
)
|
||
target_ids = torch.tensor(self.ctc_tokenizer.encode(result["text"]), dtype=torch.int64)
|
||
result["timestamps"] = forced_align(ctc_result["ctc_logits"], target_ids, self.blank_id)
|
||
for timestamps in [result["timestamps"], result["ctc_timestamps"]]:
|
||
for timestamp in timestamps:
|
||
timestamp["token"] = self.ctc_tokenizer.decode([timestamp["token"]])
|
||
timestamp["start_time"] = timestamp["start_time"] * 6 * 10 / 1000
|
||
timestamp["end_time"] = timestamp["end_time"] * 6 * 10 / 1000
|
||
|
||
if ibest_writer is not None:
|
||
ibest_writer["text"][key[0]] = response.replace("\n", " ")
|
||
ibest_writer["label"][key[0]] = label.replace("\n", " ")
|
||
ibest_writer["text_tn"][key[0]] = response_clean
|
||
|
||
return results, meta_data
|
||
|
||
@staticmethod
|
||
def from_pretrained(model: str = None, **kwargs):
|
||
"""From pretrained.
|
||
|
||
Args:
|
||
model: Model instance or model name.
|
||
**kwargs: Additional keyword arguments.
|
||
"""
|
||
from funasr import AutoModel
|
||
|
||
model, kwargs = AutoModel.build_model(model=model, trust_remote_code=True, **kwargs)
|
||
|
||
return model, kwargs
|