Files
jarvis-models/src/blackbox/asr.py
2024-08-22 15:26:27 +08:00

136 lines
4.7 KiB
Python

from io import BytesIO
from typing import Any, Coroutine
from fastapi import Request, Response, status
from fastapi.responses import JSONResponse
from ..asr.rapid_paraformer.utils import read_yaml
from ..asr.rapid_paraformer import RapidParaformer
from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess
from .blackbox import Blackbox
from injector import singleton, inject
import tempfile
import json
import os
from ..configuration import SenseVoiceConf
from ..log.logging_time import logging_time
import logging
logger = logging.getLogger(__name__)
@singleton
class ASR(Blackbox):
mode: str
url: str
speed: int
device: str
language: str
speaker: str
@logging_time(logger=logger)
def model_init(self, sensevoice_config: SenseVoiceConf) -> None:
config = read_yaml(".env.yaml")
self.paraformer = RapidParaformer(config)
model_dir = "/home/gpu/Workspace/Models/SenseVoice/SenseVoiceSmall"
self.speed = sensevoice_config.speed
self.device = sensevoice_config.device
self.language = sensevoice_config.language
self.speaker = sensevoice_config.speaker
self.device = sensevoice_config.device
self.url = ''
self.mode = sensevoice_config.mode
self.asr = None
self.speaker_ids = None
# os.environ['CUDA_VISIBLE_DEVICES'] = str(sensevoice_config.device)
if self.mode == 'local':
self.asr = AutoModel(
model=model_dir,
trust_remote_code=True,
remote_code= "/home/gpu/Workspace/SenseVoice/model.py",
vad_model="fsmn-vad",
vad_kwargs={"max_single_segment_time": 30000},
device=self.device,
)
else:
self.url = sensevoice_config.url
logging.info('#### Initializing SenseVoiceASR Service in cuda:' + sensevoice_config.device + ' mode...')
@inject
def __init__(self, sensevoice_config: SenseVoiceConf, settings: dict) -> None:
self.model_init(sensevoice_config)
def __call__(self, *args, **kwargs):
return self.processing(*args, **kwargs)
async def processing(self, *args, settings: dict):
print("\nChat Settings: ", settings)
if settings is None:
settings = {}
user_model_name = settings.get("asr_model_name")
print(f"asr_model_name: {user_model_name}")
data = args[0]
if user_model_name == 'sensevoice' or ['sensevoice']:
# 创建一个临时文件
with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_audio_file:
temp_audio_file.write(data)
temp_audio_path = temp_audio_file.name
res = self.asr.generate(
input=temp_audio_path,
cache={},
language="auto", # "zh", "en", "yue", "ja", "ko", "nospeech"
use_itn=True,
batch_size_s=60,
merge_vad=True, #
merge_length_s=15,
)
# results = self.paraformer([BytesIO(data)])
results = rich_transcription_postprocess(res[0]["text"])
os.remove(temp_audio_path)
if len(results) == 0:
return None
return results
elif user_model_name == 'funasr' or ['funasr']:
return results
else:
results = self.paraformer([BytesIO(data)])
if len(results) == 0:
return None
return results[0]
def valid(self, data: any) -> bool:
if isinstance(data, bytes):
return True
return False
async def fast_api_handler(self, request: Request) -> Response:
data = (await request.form()).get("audio")
setting: dict = (await request.form()).get("settings")
if isinstance(setting, str):
try:
setting = json.loads(setting) # 尝试将字符串转换为字典
except json.JSONDecodeError:
return JSONResponse(content={"error": "Invalid settings format"}, status_code=status.HTTP_400_BAD_REQUEST)
if data is None:
# self.logger.warn("asr bag request","type", "fast_api_handler", "api", "asr")
return JSONResponse(content={"error": "data is required"}, status_code=status.HTTP_400_BAD_REQUEST)
d = await data.read()
try:
txt = await self.processing(d, settings=setting)
except ValueError as e:
return JSONResponse(content={"error": str(e)}, status_code=status.HTTP_400_BAD_REQUEST)
return JSONResponse(content={"text": txt}, status_code=status.HTTP_200_OK)