import uvicorn
import requests
import os
from fastapi import FastAPI
from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess
from pydantic import BaseModel

model = AutoModel(
    model="iic/SenseVoiceSmall",
    trust_remote_code=True,
    remote_code="./model.py",
    vad_model="fsmn-vad",
    vad_kwargs={"max_single_segment_time": 30000},
    device="cuda:0",
)


# 定义asr数据模型，用于接收POST请求中的数据
class ASRItem(BaseModel):
    url: str  # 输入音频的远程URL


app = FastAPI()

@app.post("/asr")
async def asr(item: ASRItem):
    try:
        # 下载远程音频文件
        response = requests.get(item.url)
        if response.status_code != 200:
            raise Exception(f"Failed to download audio from {item.url}. Status code: {response.status_code}")
        with open("test.wav", "wb") as f:
            f.write(response.content)

        # 进行语音识别
        res = model.generate("test.wav",
                             language="auto",  # "zn", "en", "yue", "ja", "ko", "nospeech"
                             use_itn=True,
                             batch_size_s=60,
                             merge_vad=True,  #
                             merge_length_s=15, )
        text = rich_transcription_postprocess(res[0]["text"])
        print(text)
        result_dict = {"code": 0, "msg": "ok", "res": text}
    except Exception as e:
        result_dict = {"code": 1, "msg": str(e)}
    finally:
        # 删除临时下载的文件
        if os.path.exists("test.wav"):
            os.remove("test.wav")
    return result_dict


if __name__ == '__main__':
    uvicorn.run(app, host='0.0.0.0', port=2003)