#!/usr/bin/env python3 # -*- encoding: utf-8 -*- # Copyright FunASR (https://github.com/FunAudioLLM/SenseVoice). All Rights Reserved. # MIT License (https://opensource.org/licenses/MIT) import os.path from pathlib import Path from typing import List, Union, Tuple TORCH_AVAILABLE = False import numpy as np import axengine as axe try: import librosa except ImportError: print("Warning: librosa not found. Please install it using 'pip install librosa'.") # Provide a fallback implementation if needed def load_wav_fallback(path, sr=None): import wave import numpy as np with wave.open(path, 'rb') as wf: num_frames = wf.getnframes() frames = wf.readframes(num_frames) return np.frombuffer(frames, dtype=np.int16).astype(np.float32) / 32768.0, wf.getframerate() from utils.infer_utils import ( CharTokenizer, get_logger, read_yaml, ) from utils.frontend import WavFrontend logging = get_logger() def sequence_mask_np(lengths, maxlen=None, dtype=np.float32): lengths = np.asarray(lengths).astype(np.int64) if maxlen is None: maxlen = int(lengths.max()) row_vector = np.arange(0, maxlen, 1) matrix = lengths.reshape(-1, 1) mask = row_vector < matrix return mask.astype(dtype) def unique_consecutive_np(arr: np.ndarray) -> np.ndarray: if arr.size == 0: return arr out = [arr[0]] for v in arr[1:]: if v != out[-1]: out.append(v) return np.asarray(out, dtype=arr.dtype) class AX_SenseVoiceSmall: """ Author: Speech Lab of DAMO Academy, Alibaba Group Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition https://arxiv.org/abs/2206.08317 """ def __init__( self, model_dir: Union[str, Path] = None, batch_size: int = 1, seq_len: int = 68 ): model_file = os.path.join(model_dir, "sensevoice.axmodel") config_file = os.path.join(model_dir, "sensevoice/config.yaml") cmvn_file = os.path.join(model_dir, "sensevoice/am.mvn") config = read_yaml(config_file) self.model_dir = model_dir # token_list = os.path.join(model_dir, "tokens.json") # with open(token_list, "r", encoding="utf-8") as f: # token_list = json.load(f) # self.converter = TokenIDConverter(token_list) self.tokenizer = CharTokenizer() config["frontend_conf"]['cmvn_file'] = cmvn_file self.frontend = WavFrontend(**config["frontend_conf"]) # self.ort_infer = OrtInferSession( # model_file, device_id, intra_op_num_threads=intra_op_num_threads # ) self.session = axe.InferenceSession(model_file) self.batch_size = batch_size self.blank_id = 0 self.seq_len = seq_len 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} def __call__(self, wav_content: Union[str, np.ndarray, List[str]], language: str, withitn: bool, position_encoding: np.ndarray, tokenizer=None, **kwargs) -> List: """Enhanced model inference with additional features from model.py Args: wav_content: Audio data or path language: Language code for processing withitn: Whether to use ITN (inverse text normalization) position_encoding: Position encoding tensor tokenizer: Tokenizer for text conversion **kwargs: Additional arguments """ # Start time tracking for metadata import time meta_data = {} time_start = time.perf_counter() # Load waveform data waveform_list = self.load_data(wav_content, self.frontend.opts.frame_opts.samp_freq) waveform_nums = len(waveform_list) time_load = time.perf_counter() meta_data["load_data"] = f"{time_load - time_start:0.3f}" # Get key for result identification key = kwargs.get("key", ["wav_file"]) if isinstance(wav_content, str): wav_name = os.path.splitext(os.path.basename(wav_content))[0] if key == ["wav_file"]: key = [wav_name] # Load queries from saved numpy files language_query = np.load(os.path.join(self.model_dir, f"{language}.npy")) textnorm_query = np.load(os.path.join(self.model_dir, "withitn.npy") if withitn else os.path.join(self.model_dir, "woitn.npy")) event_emo_query = np.load(os.path.join(self.model_dir, "event_emo.npy")) # Concatenate queries to form input_query input_query = np.concatenate((language_query, event_emo_query, textnorm_query), axis=1) # Setup dataset directories for saving intermediate files dataset = "dataset" os.makedirs(dataset, exist_ok=True) speech_dir = os.path.join(dataset, "speech", language, "withitn" if withitn else "woitn") mask_dir = os.path.join(dataset, "masks", language, "withitn" if withitn else "woitn") pe_dir = os.path.join(dataset, "position_encoding", language, "withitn" if withitn else "woitn") os.makedirs(speech_dir, exist_ok=True) os.makedirs(mask_dir, exist_ok=True) os.makedirs(pe_dir, exist_ok=True) # Process features results = [] output_timestamp = kwargs.get("output_timestamp", False) ban_emo_unk = kwargs.get("ban_emo_unk", False) ibest_writer = None # 添加时间偏移变量，用于跟踪连续的时间戳 time_offset = 0.0 # 添加合并时间戳变量，用于存储同一音频文件的所有时间戳，仅保留必要功能 merged_timestamps = [] merged_words = [] # Handle output_dir without using DatadirWriter (which is not available) output_dir = kwargs.get("output_dir") slice_len = self.seq_len - 4 time_pre = time.perf_counter() meta_data["preprocess"] = f"{time_pre - time_load:0.3f}" for beg_idx in range(0, waveform_nums, self.batch_size): end_idx = min(waveform_nums, beg_idx + self.batch_size) feats, feats_len = self.extract_feat(waveform_list[beg_idx:end_idx]) time_feat = time.perf_counter() meta_data["extract_feat"] = f"{time_feat - time_pre:0.3f}" for i in range(int(np.ceil(feats.shape[1] / slice_len))): sub_feats = np.concatenate([input_query, feats[:, i*slice_len : (i+1)*slice_len, :]], axis=1) feats_len[0] = sub_feats.shape[1] # 计算当前片段的实际长度（帧数） actual_slice_length = min(slice_len, feats.shape[1] - i*slice_len) if feats_len[0] < self.seq_len: sub_feats = np.concatenate([sub_feats, np.zeros((1, self.seq_len - feats_len[0], 560), dtype=np.float32)], axis=1) masks = sequence_mask_np([self.seq_len], maxlen=self.seq_len, dtype=np.float32)[:, None, :] # Run inference ctc_logits, encoder_out_lens = self.infer(sub_feats, masks, position_encoding) # Ban emotion unknown token if requested if ban_emo_unk: ctc_logits[:, :, self.emo_dict["unk"]] = -float("inf") # Process results for each batch b, n, d = ctc_logits.shape if isinstance(key, (list, tuple)) and len(key) < b: key = key * b for j in range(b): enc_len = int(encoder_out_lens[j]) if not hasattr(encoder_out_lens[j], "item") else int(encoder_out_lens[j].item()) x = ctc_logits[j, : enc_len, :] yseq = np.argmax(x, axis=-1) yseq = unique_consecutive_np(yseq) mask = yseq != self.blank_id token_int = yseq[mask].tolist() # Convert tokens to text text = tokenizer.decode(token_int) if tokenizer is not None else str(token_int) # 文本处理 # 简化文本处理，不执行合并操作 # Write to output directory if provided if ibest_writer is not None: ibest_writer["text"][key[j]] = text if output_timestamp: # Torch-free build: skip timestamp generation output_timestamp = False result_i = {"key": key[j] if j < len(key) else f"result_{j}", "text": text} # 直接添加结果，重复处理将在export.py中进行 results.append(result_i) time_offset = round(time_offset + (actual_slice_length * 60 ) / 1000, 2) # 结果处理完成 time_end = time.perf_counter() meta_data["total_time"] = f"{time_end - time_start:0.3f}" meta_data["inference_time"] = f"{time_end - time_feat:0.3f}" # 简化元数据处理，只保留时间戳 if len(results) > 0 and output_timestamp and merged_timestamps: # 只保留时间戳数据，以便在export.py中处理 meta_data["merged_timestamps"] = merged_timestamps meta_data["merged_words"] = merged_words return results, meta_data def load_data(self, wav_content: Union[str, np.ndarray, List[str]], fs: int = None) -> List: def load_wav(path: str) -> np.ndarray: try: # Use librosa if available if 'librosa' in globals(): waveform, _ = librosa.load(path, sr=fs) else: # Use fallback implementation waveform, native_sr = load_wav_fallback(path) if fs is not None and native_sr != fs: # Implement resampling if needed print(f"Warning: Resampling from {native_sr} to {fs} is not implemented in fallback mode") return waveform except Exception as e: print(f"Error loading audio file {path}: {e}") # Return empty audio in case of error return np.zeros(1600, dtype=np.float32) if isinstance(wav_content, np.ndarray): return [wav_content] if isinstance(wav_content, str): return [load_wav(wav_content)] if isinstance(wav_content, list): return [load_wav(path) for path in wav_content] raise TypeError(f"The type of {wav_content} is not in [str, np.ndarray, list]") def extract_feat(self, waveform_list: List[np.ndarray]) -> Tuple[np.ndarray, np.ndarray]: feats, feats_len = [], [] for waveform in waveform_list: speech, _ = self.frontend.fbank(waveform) 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 @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 infer(self, feats: np.ndarray, masks: np.ndarray, position_encoding: np.ndarray, ) -> Tuple[np.ndarray, np.ndarray]: #outputs = self.ort_infer([feats, masks, position_encoding]) outputs =self.session.run(None, { 'speech': feats, 'masks': masks, 'position_encoding': position_encoding }) return outputs