import axengine as axe import numpy as np import librosa from frontend import WavFrontend import os import time from typing import Any, Dict, Iterable, List, NamedTuple, Set, Tuple, Union from print_utils import rich_transcription_postprocess def sequence_mask(lengths, maxlen=None, dtype=np.float32): # 如果 maxlen 未指定，则取 lengths 中的最大值 if maxlen is None: maxlen = np.max(lengths) # 创建一个从 0 到 maxlen-1 的行向量 row_vector = np.arange(0, maxlen, 1) # 将 lengths 转换为列向量 matrix = np.expand_dims(lengths, axis=-1) # 比较生成掩码 mask = row_vector < matrix if mask.shape[-1] < lengths[0]: mask = np.concatenate([mask, np.zeros((mask.shape[0], lengths[0] - mask.shape[-1]), dtype=np.float32)], axis=-1) # 返回指定数据类型的掩码 return mask.astype(dtype)[None, ...] def unique_consecutive_np(x, dim=None, return_inverse=False, return_counts=False): if dim is None: # 默认情况，展平后去重 x_flat = x.ravel() mask = np.concatenate(([True], x_flat[1:] != x_flat[:-1])) unique_data = x_flat[mask] else: # 沿着指定维度去重 axis = dim if dim >= 0 else x.ndim + dim if axis >= x.ndim: raise ValueError(f"dim {dim} is out of range for array of dimension {x.ndim}") # 使用 np.diff 检查相邻元素是否相同 mask = np.ones(x.shape[axis], dtype=bool) if x.shape[axis] > 1: # 比较当前元素和前一个元素是否不同 diff = np.diff(x, axis=axis) mask[1:] = np.any(diff != 0, axis=tuple(range(diff.ndim))[axis:]) # 使用 mask 索引提取唯一元素 unique_data = np.take(x, np.where(mask)[0], axis=axis) # 处理 return_inverse 和 return_counts results = (unique_data,) if return_inverse: if dim is None: inv_idx = np.cumsum(mask) - 1 else: inv_idx = np.cumsum(mask) - 1 # 需要调整形状以匹配输入 inv_idx = np.expand_dims(inv_idx, axis=axis) inv_idx = np.broadcast_to(inv_idx, x.shape) results += (inv_idx,) if return_counts: if dim is None: counts = np.diff(np.where(np.concatenate((mask, [True])))[0]) else: counts = np.diff(np.where(np.concatenate((mask, [True])))[0]) results += (counts,) return results[0] if len(results) == 1 else results def longest_common_suffix_prefix_with_tolerance( lhs, rhs, tolerate: int = 0 ) -> int: """ 计算两个数组的最长公共子序列，该子序列必须同时满足： - 是 lhs 的后 n 个元素（后缀） - 是 rhs 的前 n 个元素（前缀） 并且允许最多 `tolerate` 个元素不匹配。 参数: lhs: np.ndarray, 第一个数组 rhs: np.ndarray, 第二个数组 tolerate: int, 允许的不匹配元素数量（默认为 0，即完全匹配） 返回: int: 最长公共后缀/前缀的长度（如果没有则返回 0） """ max_possible_n = min(len(lhs), len(rhs)) for n in range(max_possible_n, 0, -1): mismatches = np.sum(lhs[-n:] != rhs[:n]) if mismatches <= tolerate: return n return 0 class SenseVoiceAx: def __init__(self, model_path, max_len=68, language="auto", use_itn=True, tokenizer=None): model_path_root = os.path.join(os.path.dirname(model_path), "..") embedding_root = os.path.join(model_path_root, "embeddings") self.frontend = WavFrontend(cmvn_file=f"{model_path_root}/am.mvn", fs=16000, window="hamming", n_mels=80, frame_length=25, frame_shift=10, lfr_m=7, lfr_n=6,) self.model = axe.InferenceSession(model_path) self.sample_rate = 16000 self.tokenizer = tokenizer self.blank_id = 0 self.max_len = max_len self.padding = 16 self.lid_dict = {"auto": 0, "zh": 3, "en": 4, "yue": 7, "ja": 11, "ko": 12, "nospeech": 13} self.lid_int_dict = {24884: 3, 24885: 4, 24888: 7, 24892: 11, 24896: 12, 24992: 13} self.textnorm_dict = {"withitn": 14, "woitn": 15} self.textnorm_int_dict = {25016: 14, 25017: 15} self.emo_dict = {"unk": 25009, "happy": 25001, "sad": 25002, "angry": 25003, "neutral": 25004} self.position_encoding = np.load(f"{embedding_root}/position_encoding.npy") language_query = np.load(f"{embedding_root}/{language}.npy") textnorm_query = np.load(f"{embedding_root}/withitn.npy") if use_itn else np.load(f"{embedding_root}/woitn.npy") event_emo_query = np.load(f"{embedding_root}/event_emo.npy") self.input_query = np.concatenate((textnorm_query, language_query, event_emo_query), axis=1) self.query_num = self.input_query.shape[1] def load_data(self, filepath: str) -> np.ndarray: waveform, _ = librosa.load(filepath, sr=self.sample_rate) return waveform.flatten() @staticmethod def pad_feats(feats: List[np.ndarray], max_feat_len: int) -> np.ndarray: def pad_feat(feat: np.ndarray, cur_len: int) -> np.ndarray: pad_width = ((0, max_feat_len - cur_len), (0, 0)) return np.pad(feat, pad_width, "constant", constant_values=0) feat_res = [pad_feat(feat, feat.shape[0]) for feat in feats] feats = np.array(feat_res).astype(np.float32) return feats def preprocess(self, waveform): feats, feats_len = [], [] for wf in [waveform]: speech, _ = self.frontend.fbank(wf) feat, feat_len = self.frontend.lfr_cmvn(speech) feats.append(feat) feats_len.append(feat_len) feats = self.pad_feats(feats, np.max(feats_len)) feats_len = np.array(feats_len).astype(np.int32) return feats, feats_len def postprocess(self, ctc_logits, encoder_out_lens): # 提取数据 x = ctc_logits[0, :encoder_out_lens[0], :] # 获取最大值索引 yseq = np.argmax(x, axis=-1) # 去除连续重复元素 yseq = unique_consecutive_np(yseq, dim=-1) # 创建掩码并过滤 blank_id mask = yseq != self.blank_id token_int = yseq[mask].tolist() return token_int def infer_waveform(self, waveform: np.ndarray): feat, feat_len = self.preprocess(waveform) slice_len = self.max_len - self.query_num slice_num = int(np.ceil(feat.shape[1] / slice_len)) asr_res = [] prev_token_int = None for i in range(slice_num): if i == 0: sub_feat = feat[:, i*slice_len:(i+1)*slice_len, :] else: sub_feat = feat[:, i*slice_len - self.padding:(i+1)*slice_len - self.padding, :] # concat query sub_feat = np.concatenate([self.input_query, sub_feat], axis=1) real_len = sub_feat.shape[1] if real_len < self.max_len: sub_feat = np.concatenate([ sub_feat, np.zeros((1, self.max_len - real_len, sub_feat.shape[-1]), dtype=np.float32) ], axis=1) masks = sequence_mask(np.array([self.max_len], dtype=np.int32), maxlen=real_len, dtype=np.float32) outputs = self.model.run(None, {"speech": sub_feat, "masks": masks, "position_encoding": self.position_encoding}) ctc_logits, encoder_out_lens = outputs token_int = self.postprocess(ctc_logits, encoder_out_lens) # common prefix if self.padding > 0 and prev_token_int is not None: # prefix_len = common_prefix_len(prev_token_int, token_int) prefix_len = longest_common_suffix_prefix_with_tolerance(prev_token_int, token_int, 6) common_prefix = rich_transcription_postprocess(self.tokenizer.tokens2text(token_int[:prefix_len])) asr_res[-1] = asr_res[-1][:-len(common_prefix)] prev_token_int = np.copy(token_int) asr_res.append(self.tokenizer.tokens2text(token_int)) return asr_res def infer(self, filepath_or_data: Union[np.ndarray, str], print_rtf=True): if isinstance(filepath_or_data, str): waveform = self.load_data(filepath_or_data) else: waveform = filepath_or_data total_time = waveform.shape[-1] / self.sample_rate start = time.time() asr_res = self.infer_waveform(waveform) latency = time.time() - start if print_rtf: rtf = latency / total_time print(f"RTF: {rtf} Latency: {latency}s Total length: {total_time}s") return asr_res