SenseVoiceAx.py

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

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
    
    # 返回指定数据类型的掩码
    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

class SenseVoiceAx:
    def __init__(self, model_path, language="auto", use_itn=True, tokenizer=None):
        model_path_root = os.path.join(os.path.dirname(model_path), "../embeddings")
        self.frontend = WavFrontend(cmvn_file="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 = 34

        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"{model_path_root}/position_encoding.npy")
        language_query = np.load(f"{model_path_root}/{language}.npy")
        textnorm_query = np.load(f"{model_path_root}/withitn.npy") if use_itn else np.load(f"{model_path_root}/woitn.npy")
        event_emo_query = np.load(f"{model_path_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]
        self.masks = sequence_mask(np.array([self.max_len], dtype=np.int32), dtype=np.float32)

    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 = []
        for i in range(slice_num):
            sub_feat = feat[:, i*slice_len:(i+1)*slice_len, :]
            # concat query
            sub_feat = np.concatenate([self.input_query, sub_feat], axis=1)

            if sub_feat.shape[1] < self.max_len:
                sub_feat = np.concatenate([
                        sub_feat, 
                        np.zeros((1, self.max_len - sub_feat.shape[1], sub_feat.shape[-1]), dtype=np.float32)
                    ],
                    axis=1)
                
            outputs = self.model.run(None, {"speech": sub_feat,
                                            "masks": self.masks,
                                            "position_encoding": self.position_encoding})
            ctc_logits, encoder_out_lens = outputs

            token_int = self.postprocess(ctc_logits, encoder_out_lens)
            if self.tokenizer is not None:
                asr_res.append(self.tokenizer.tokens2text(token_int))
            else:
                asr_res.append(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