demo.py

#!/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
import torch
import argparse
from model import  SinusoidalPositionEncoder
from utils.ax_model_bin import AX_SenseVoiceSmall
from utils.ax_vad_bin import AX_Fsmn_vad 
from utils.vad_utils import merge_vad
from utils.ax_cam_bin import AX_SpeakerEmbeddingInference, do_clustering, distribute_spk, get_trans_sentence_sensevoice
from funasr.tokenizer.sentencepiece_tokenizer import SentencepiecesTokenizer
from utils.ax_cam_bin import chunk
import time
import librosa
import soundfile as sf
import numpy as np
from utils.infer_func import InferManager
from ml_dtypes import bfloat16
from transformers import AutoConfig, AutoTokenizer
from loguru import logger


llm_axmodel_path = "./ax_model/Qwen3-4B-Instruct-2507-GPTQ-Int4_8k_axmodel"
llm_hf_tokenizer_path = "./tokenizer_qwen3_int4"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
embeds = np.load(os.path.join(llm_axmodel_path, "model.embed_tokens.weight.npy"))


def parse_args():
    parser = argparse.ArgumentParser(description="SenseVoice inference script")
    parser.add_argument("--output_dir", type=str, default="./output_dir", help="Output directory")
    parser.add_argument("--seq_len", type=int, default=132, help="Sequence length for model") #68 ,132
    parser.add_argument("--wav_file", type=str, default="wav/vad_example.wav",help="Input wav file")
    return parser.parse_args()


def run_model(prompt, max_seq_len=8191, slice_len=256, max_prefill_len=4095):
    messages = [
        {
            "role": "system",
            "content": "你叫小惠, 你是一个专业的会议记录分析助手, 善于从会议记录(按照时间先后记录不同人物的发言)中提取关键信息并生成合适的总结. \n 请你基于以下会议记录, 在深度思考后, 总结这段会议记录的参会人员以及内容摘要. ",
        },
        {
            "role": "user",
            "content": prompt,
        },
    ]

    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(device)
    input_ids = model_inputs.input_ids

    ######################################################################

    token_ids = input_ids[0].cpu().numpy().tolist()
    token_len = len(token_ids)

    assert token_len <= max_prefill_len, f"Input token length {token_len} exceeds max prefill length {max_prefill_len}"
    # import pdb; pdb.set_trace()
    prefill_data = np.take(embeds, token_ids, axis=0)
    prefill_data = prefill_data.astype(bfloat16)

    imer = InferManager(cfg, llm_axmodel_path, max_seq_len=max_seq_len, max_prefill_len=max_prefill_len) # prefill + decode max length
    token_ids = imer.prefill(tokenizer, token_ids, prefill_data, slice_len=slice_len)
    imer.decode(tokenizer, token_ids, embeds, slice_len=slice_len, eos_token_id=eos_token_id)

    # 分别打印输入和输出 token 数, kv cache 总长度
    print(f"\n输入 token 数: {token_len}")
    print(f"输出 token 数: {len(token_ids) - token_len}")
    print(f"kv cache 总长度: {max_seq_len}")

    # 保存输出 token
    output_text = tokenizer.decode(token_ids[token_len:], skip_special_tokens=True)
    return output_text

if __name__ == "__main__":
    args = parse_args()
    seq_len = args.seq_len
    model_path = args.output_dir
    os.makedirs(model_path, exist_ok=True)

    print(f"Initializing  model...")

    ax_model_dir = "ax_model"
    total_inference_start = time.time()
   
    model_vad = AX_Fsmn_vad(ax_model_dir)
    speaker_model = AX_SpeakerEmbeddingInference(model_dir=ax_model_dir)
    embed = SinusoidalPositionEncoder()
    position_encoding = embed.get_position_encoding(torch.randn(1, seq_len, 560)).numpy()
    model_bin = AX_SenseVoiceSmall(ax_model_dir, seq_len=seq_len)

    # build tokenizer
    print(f"Loading tokenizer...")
    tokenizer = None
    tokenizer_path = os.path.join(ax_model_dir, "chn_jpn_yue_eng_ko_spectok.bpe.model")
    tokenizer = SentencepiecesTokenizer(bpemodel=tokenizer_path)

    # Set up audio file for processing
    wav_file = args.wav_file #S_R004S03C01

    print(f"Running inference on example file...")
    
    withitn = True
    norm_type = "withitn" if withitn else "woitn"
    print(f"\nProcessing with text normalization: {norm_type}")
        
    language = "auto"
    print(f"\n--- Processing language: {language} ---")
    print(f"Processing file: {wav_file}")
    inference_start = time.time()
    # 加载音频数据
    
    speech, fs = librosa.load(wav_file, sr=None)
    # 检查采样率，如果不是16kHz则进行重采样
    if fs != 16000:
        print(f"Resampling audio from {fs}Hz to 16000Hz")
        speech = librosa.resample(y=speech, orig_sr=fs, target_sr=16000)
        fs = 16000
    audio_duration = librosa.get_duration(y=speech, sr=fs)
    speech_lengths = len(speech)

    try:
        #增加vad model 推理及处理
        vad_start_time = time.time()
        res_vad = model_vad(speech)[0]
        vad_segments = merge_vad(res_vad, 15 * 1000)  #短语音段合并 # vad_segments: [[0, 6480], [6480, 23670], [23670, 38210], [38210, 49910], [49910, 59820], [59820, 70550]]
        vad_time_cost = time.time() - vad_start_time
        print(f"VAD processing time: {vad_time_cost:.2f} seconds")
        
        # emb_extraction
        vad_time = [[vad_t[0]/1000, vad_t[1]/1000] for vad_t in res_vad]
        chunks = [c for (st, ed) in vad_time for c in chunk(st, ed)]

        # Extract speaker embeddings for each chunk
        print("Extracting speaker embeddings...")
        speaker_start_time = time.time()
        # import pdb 
        # pdb.set_trace()
        embeddings = speaker_model(speech, fs, chunks=chunks)
        speaker_time_cost = time.time() - speaker_start_time
        print(f"Speaker embedding extraction time: {speaker_time_cost:.2f} seconds")
        print(f"Generated embeddings shape: {embeddings.shape}")

        clustering_start_time = time.time()
        speaker_num, diar_results = do_clustering(chunks, embeddings, speaker_num=None)
        clustering_time_cost = time.time() - clustering_start_time
        print(f"Speaker clustering time: {clustering_time_cost:.2f} seconds")
        
        # print(f"VAD segments detected: {len(vad_segments)}")
        
        # 存储所有分片结果
        all_results = []
        all_metadata = {}
        
        # 遍历每个VAD片段并处理
        asr_start_time = time.time()
        for i, segment in enumerate(vad_segments):
            segment_start, segment_end = segment
            # 从原始音频中提取该片段
            start_sample = int(segment_start / 1000 * fs)
            end_sample = min(int(segment_end / 1000 * fs), speech_lengths)
            segment_speech = speech[start_sample:end_sample]
            
            # 计算时间偏移量（毫秒转秒）
            time_offset_sec = segment_start / 1000.0
            
            # 为当前片段创建临时文件
            segment_filename = f"temp_segment_{i}.wav"
            sf.write(segment_filename, segment_speech, fs)
            
            # 对当前片段进行识别
            try:
                segment_res, segment_meta = model_bin(
                    segment_filename, 
                    language, 
                    withitn, 
                    position_encoding, 
                    tokenizer=tokenizer,
                    output_timestamp=True,
                    ban_emo_unk=False,
                    output_dir=model_path,
                    key=[f"{os.path.basename(wav_file)}_segment_{i}"]
                )

                if "merged_words" in segment_meta:
                    if "merged_words" not in all_metadata:
                        all_metadata["merged_words"] = []
                    all_metadata["merged_words"].extend(segment_meta["merged_words"])
                    
                if "merged_timestamps" in segment_meta:
                    if "merged_timestamps" not in all_metadata:
                        all_metadata["merged_timestamps"] = []
                    adjusted_timestamps = [[min(ts[0] + time_offset_sec, audio_duration), min(ts[1] + time_offset_sec, audio_duration)]  
                                        for ts in segment_meta["merged_timestamps"]]  # 确保不超过音频总时长
                    all_metadata["merged_timestamps"].extend(adjusted_timestamps)
                
                if os.path.exists(segment_filename):
                    os.remove(segment_filename)
                    
            except Exception as e:
                if os.path.exists(segment_filename):
                    os.remove(segment_filename)
                raise
    
        output_asr = {
            "merged_words": all_metadata.get("merged_words", []),
            "merged_timestamps": all_metadata.get("merged_timestamps", [])
        }

    except Exception as e:
        raise

    asr_time_cost = time.time() - asr_start_time
    print(f"ASR processing time: {asr_time_cost:.2f} seconds")

    asr_timestamps = get_trans_sentence_sensevoice(output_asr)
    sentence_info_with_spk = distribute_spk(asr_timestamps, diar_results)
    inference_time_cost = time.time() - inference_start
    inference_time_cost_all = time.time() - total_inference_start
    rtf = inference_time_cost / audio_duration
    print(f"Inference time for {wav_file}: {inference_time_cost:.2f} seconds")
    print(f"load model  + Inference time for {wav_file}: {inference_time_cost_all:.2f} seconds")
    print(f"Audio duration: {audio_duration:.2f} seconds")
    print(f"RTF: {rtf:.2f}")

    # Save Results
    output_trans_path = os.path.join(args.output_dir, f"{wav_file.split('/')[-1]}.txt")
    
    try:
        with open(output_trans_path, 'w', encoding='utf-8') as f:
            for text_string, timeinterval, spk in sentence_info_with_spk:
                f.write(f'Speaker_{spk}: [{timeinterval[0]:.3f} {timeinterval[1]:.3f}] {text_string}\n')
    except Exception as e:
        pass
    ########################### LLM Inference #############################
    logger.info("Starting LLM Inference for meeting summary...")
    # load the tokenizer and the model
    tokenizer = AutoTokenizer.from_pretrained(llm_hf_tokenizer_path)
    cfg = AutoConfig.from_pretrained(llm_hf_tokenizer_path, trust_remote_code=True)

    eos_token_id = None
    if isinstance(cfg.eos_token_id, list) and len(cfg.eos_token_id) > 1:
        eos_token_id = cfg.eos_token_id

    with open(output_trans_path, 'r', encoding='utf-8') as file:
        prompt = file.read()

    prompt_list = prompt.split('\n')

    output_text_l = []
    cur_prompt = ""
    for idx, prompt in enumerate(prompt_list):

        cur_prompt += prompt + "\n"

        if len(cur_prompt) < 5000 and idx != len(prompt_list) -1:
            continue

        output_text = run_model(cur_prompt)
        output_text_l.append(output_text)
        print("\n")

        cur_prompt = ""
    
    # TODO: 如果输入文本很长, 可以考虑再次调用模型对分段总结进行全面总结, 但是此时也需要保证总输入长度小于模型规定的上限
    # print("\n\n开始进行最终总结...")
    # prompt = "\n".join(output_text_l)
    # final_output_text = run_model(prompt)
    # print("最终总结结果:\n", final_output_text)

    # 将 output_text_l 保存为 md 格式的文件, 按照原始 LLM 的输出格式进行保存
    output_summary_path = os.path.join(args.output_dir, f"{wav_file.split('/')[-1]}_summary.md")
    output_text_combined = "\n\n".join(output_text_l)
    with open(output_summary_path, 'w', encoding='utf-8') as f:
        f.write(output_text_combined)

    logger.info(f"LLM Inference completed. Summary saved to {output_summary_path}")