from funasr import AutoModel

def transcribe_audio(audio_path):
    """
    基于funasr的sensevoicesmall模型进行带时间戳的音频转录
    
    参数:
    audio_path (str): 输入音频文件的路径
    
    返回:
    list: 包含转录文本及其时间戳的列表，每个元素是一个字典，包含 'text', 'start', 'end' 字段
    """
    # 初始化模型，使用 sensevoicesmall 进行语音识别，开启时间戳功能
    model = AutoModel(model="iic/speech_paraformer-large-vad-punc_asr_nat-zh-cn-16k-common-vocab8404-pytorch")
    # model = AutoModel(model="iic/sensevoicesmall")
    
    # 调用模型进行流式识别
    result = model.generate(input=audio_path, timestamp=True)
    print(result)
    
    # 解析结果，提取带时间戳的信息
    timestamped_transcriptions = []
    if result:
        for item in result[0]['st']:
            timestamped_transcriptions.append({
                'text': item['text'],
                'start': item['start'],
                'end': item['end']
            })
    
    return timestamped_transcriptions


if __name__ == "__main__":
    audio_path = r'C:\Users\User\Desktop\Praasper\data\test_audio.wav'  # 替换为你的音频文件路径
    result = transcribe_audio(audio_path)
    # print(result)
