278 lines
11 KiB
Python
278 lines
11 KiB
Python
|
import asyncio
|
|||
|
|
import json
|
|||
|
|
import os
|
|||
|
|
import time
|
|||
|
|
import traceback
|
|||
|
|
from copy import deepcopy
|
|||
|
|
from enum import Enum
|
|||
|
|
from typing import Any, List
|
|||
|
|
|
|||
|
|
import websockets
|
|||
|
|
from pydantic_core import ValidationError
|
|||
|
|
|
|||
|
|
from schemas.context import ASRContext
|
|||
|
|
from schemas.stream import StreamResultModel, StreamWordsModel
|
|||
|
|
from utils.logger import logger
|
|||
|
|
|
|||
|
|
IN_TEST = os.getenv("SUBMIT_CONFIG_FILEPATH", None) is None
|
|||
|
|
|
|||
|
|
|
|||
|
|
class STATUS_DATA(str, Enum):
|
|||
|
|
WAITING_FIRST_INIT = "waiting_first_init"
|
|||
|
|
FIRST_FAIL = "fail"
|
|||
|
|
WAITING_SECOND_INIT = "waiting_second_init"
|
|||
|
|
SECOND_INIT = "second_fail"
|
|||
|
|
WAITING_THIRD_INIT = "waiting_third_init"
|
|||
|
|
THIRD_INIT = "third_fail"
|
|||
|
|
SUCCESS = "success"
|
|||
|
|
CLOSED = "closed"
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ClientAsync:
|
|||
|
|
def __init__(self, sut_url: str, context: ASRContext, idx: int) -> None:
|
|||
|
|
# base_url = "ws://127.0.0.1:5003"
|
|||
|
|
self.base_url = sut_url + "/recognition"
|
|||
|
|
self.context: ASRContext = deepcopy(context)
|
|||
|
|
self.idx = idx
|
|||
|
|
# if not os.getenv("DATASET_FILEPATH", ""):
|
|||
|
|
# self.base_url = "wss://speech.4paradigm.com/aibuds/api/v1/recognition"
|
|||
|
|
# self.base_url = "ws://localhost:5003/recognition"
|
|||
|
|
self.fail_count = 0
|
|||
|
|
self.close_time = 10**50
|
|||
|
|
self.send_time: List[float] = []
|
|||
|
|
self.recv_time: List[float] = []
|
|||
|
|
self.predict_data: List[Any] = []
|
|||
|
|
|
|||
|
|
async def _sender(
|
|||
|
|
self, websocket: websockets.WebSocketClientProtocol, send_queue: asyncio.Queue, recv_queue: asyncio.Queue
|
|||
|
|
):
|
|||
|
|
# 设置 websocket 缓冲区大小
|
|||
|
|
websocket.transport.set_write_buffer_limits(1024 * 1024 * 1024)
|
|||
|
|
|
|||
|
|
# 发送初始化数据
|
|||
|
|
await websocket.send(json.dumps(self._gen_init_data()))
|
|||
|
|
await send_queue.put(STATUS_DATA.WAITING_FIRST_INIT)
|
|||
|
|
connect_status = await recv_queue.get()
|
|||
|
|
if connect_status == STATUS_DATA.FIRST_FAIL:
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# 开始发送音频
|
|||
|
|
with open(self.context.file_path, "rb") as fp:
|
|||
|
|
wav_data = fp.read()
|
|||
|
|
meta_length = wav_data.index(b"data") + 8
|
|||
|
|
try:
|
|||
|
|
with open(self.context.file_path, "rb") as fp:
|
|||
|
|
# 去掉 wav 文件的头信息
|
|||
|
|
fp.read(meta_length)
|
|||
|
|
wav_time = 0.0
|
|||
|
|
label_id = 0
|
|||
|
|
char_contains_rate_checktime = []
|
|||
|
|
char_contains_rate_checktime_id = 0
|
|||
|
|
while True:
|
|||
|
|
now_time = time.perf_counter()
|
|||
|
|
chunk = fp.read(int(self.context.chunk_size))
|
|||
|
|
if not chunk:
|
|||
|
|
break
|
|||
|
|
wav_time += self.context.wait_time
|
|||
|
|
try:
|
|||
|
|
self.send_time.append(time.perf_counter())
|
|||
|
|
await asyncio.wait_for(websocket.send(chunk), timeout=0.08)
|
|||
|
|
except asyncio.exceptions.TimeoutError:
|
|||
|
|
pass
|
|||
|
|
while label_id < len(self.context.labels) and wav_time >= self.context.labels[label_id].start:
|
|||
|
|
char_contains_rate_checktime.append(now_time + 3.0)
|
|||
|
|
label_id += 1
|
|||
|
|
predict_text_len = sum(map(lambda x: len(x.text), self.predict_data))
|
|||
|
|
while char_contains_rate_checktime_id < len(char_contains_rate_checktime) and \
|
|||
|
|
char_contains_rate_checktime[char_contains_rate_checktime_id] <= now_time:
|
|||
|
|
label_text_len = sum(
|
|||
|
|
map(lambda x: len(x.answer),
|
|||
|
|
self.context.labels[:char_contains_rate_checktime_id+1]))
|
|||
|
|
if predict_text_len / self.context.char_contains_rate < label_text_len:
|
|||
|
|
self.context.fail_char_contains_rate_num += 1
|
|||
|
|
char_contains_rate_checktime_id += 1
|
|||
|
|
await asyncio.sleep(max(0, self.context.wait_time - (time.perf_counter() - now_time)))
|
|||
|
|
await websocket.send(json.dumps({"end": True}))
|
|||
|
|
logger.info(f"第 {self.idx} 条数据，当条语音数据发送完成")
|
|||
|
|
logger.info(f"第 {self.idx} 条数据，3s 后关闭双向连接.")
|
|||
|
|
self.close_time = time.perf_counter() + 3
|
|||
|
|
except websockets.exceptions.ConnectionClosedError:
|
|||
|
|
logger.error(f"第 {self.idx} 条数据发送过程中，连接断开")
|
|||
|
|
except Exception:
|
|||
|
|
logger.error(f"{traceback.print_exc()}")
|
|||
|
|
logger.error(f"第 {self.idx} 条数据，发送数据失败")
|
|||
|
|
|
|||
|
|
async def _recv(
|
|||
|
|
self, websocket: websockets.WebSocketClientProtocol, send_queue: asyncio.Queue, recv_queue: asyncio.Queue
|
|||
|
|
):
|
|||
|
|
await recv_queue.get()
|
|||
|
|
try:
|
|||
|
|
await asyncio.wait_for(websocket.recv(), timeout=2)
|
|||
|
|
except asyncio.exceptions.TimeoutError:
|
|||
|
|
await send_queue.put(STATUS_DATA.FIRST_FAIL)
|
|||
|
|
logger.info(f"第 {self.idx} 条数据，初始化阶段, 2s 没收到 success 返回，超时了")
|
|||
|
|
self.fail_count += 1
|
|||
|
|
return
|
|||
|
|
except Exception as e:
|
|||
|
|
await send_queue.put(STATUS_DATA.FIRST_FAIL)
|
|||
|
|
logger.error(f"第 {self.idx} 条数据，初始化阶段, 收到异常：{e}")
|
|||
|
|
self.fail_count += 1
|
|||
|
|
return
|
|||
|
|
else:
|
|||
|
|
await send_queue.put(STATUS_DATA.SUCCESS)
|
|||
|
|
|
|||
|
|
# 开始接收语音识别结果
|
|||
|
|
try:
|
|||
|
|
while websocket.open:
|
|||
|
|
# 接收数据
|
|||
|
|
recv_data = await websocket.recv()
|
|||
|
|
if isinstance(recv_data, str):
|
|||
|
|
self.recv_time.append(time.perf_counter())
|
|||
|
|
recv_data = str(recv_data)
|
|||
|
|
recv_data = json.loads(recv_data)
|
|||
|
|
result = StreamResultModel(**recv_data)
|
|||
|
|
recognition_results = result.asr_results
|
|||
|
|
if (
|
|||
|
|
recognition_results.final_result
|
|||
|
|
and not recognition_results.language
|
|||
|
|
and recognition_results.start_time == 0
|
|||
|
|
and recognition_results.end_time == 0
|
|||
|
|
and recognition_results.para_seq == 0
|
|||
|
|
):
|
|||
|
|
pass
|
|||
|
|
else:
|
|||
|
|
self.predict_data.append(recognition_results)
|
|||
|
|
else:
|
|||
|
|
raise Exception("返回的结果不是字符串形式")
|
|||
|
|
except websockets.exceptions.ConnectionClosedOK:
|
|||
|
|
pass
|
|||
|
|
except websockets.exceptions.ConnectionClosedError:
|
|||
|
|
pass
|
|||
|
|
except ValidationError as e:
|
|||
|
|
logger.error(f"第 {self.idx} 条数据，返回的结果不符合格式")
|
|||
|
|
logger.error(f"Exception is {e}")
|
|||
|
|
os._exit(1)
|
|||
|
|
except OSError:
|
|||
|
|
pass
|
|||
|
|
except Exception:
|
|||
|
|
logger.error(f"{traceback.print_exc()}")
|
|||
|
|
logger.error(f"第 {self.idx} 条数据，处理被测服务返回数据时出错")
|
|||
|
|
|
|||
|
|
async def _action(self):
|
|||
|
|
logger.info(f"第 {self.idx} 条数据开始测试")
|
|||
|
|
|
|||
|
|
while self.fail_count < 3:
|
|||
|
|
|
|||
|
|
send_queue = asyncio.Queue()
|
|||
|
|
recv_queue = asyncio.Queue()
|
|||
|
|
|
|||
|
|
self.send_time: List[float] = []
|
|||
|
|
self.recv_time: List[float] = []
|
|||
|
|
self.predict_data: List[Any] = []
|
|||
|
|
|
|||
|
|
async with websockets.connect(self.base_url) as websocket:
|
|||
|
|
send_task = asyncio.create_task(self._sender(websocket, send_queue, recv_queue))
|
|||
|
|
recv_task = asyncio.create_task(self._recv(websocket, recv_queue, send_queue))
|
|||
|
|
|
|||
|
|
await asyncio.gather(send_task)
|
|||
|
|
await asyncio.sleep(3)
|
|||
|
|
|
|||
|
|
await asyncio.gather(recv_task)
|
|||
|
|
|
|||
|
|
if self.send_time:
|
|||
|
|
break
|
|||
|
|
else:
|
|||
|
|
self.fail_count += 1
|
|||
|
|
logger.info(f"第 {self.idx} 条数据，初始化阶段, 第 {self.fail_count} 次失败, 1s 后重试")
|
|||
|
|
time.sleep(1)
|
|||
|
|
|
|||
|
|
def action(self):
|
|||
|
|
asyncio.run(self._action())
|
|||
|
|
return self._gen_result()
|
|||
|
|
|
|||
|
|
def _gen_result(self) -> ASRContext:
|
|||
|
|
if not self.predict_data:
|
|||
|
|
logger.error(f"第 {self.idx} 条数据，没有任何数据返回")
|
|||
|
|
self.context.append_preds(self.predict_data, self.send_time, self.recv_time)
|
|||
|
|
self.context.fail = not self.predict_data
|
|||
|
|
|
|||
|
|
punctuation_words: List[StreamWordsModel] = []
|
|||
|
|
for pred in self.predict_data:
|
|||
|
|
punctuations = [",", ".", "!", "?"]
|
|||
|
|
if pred.language == "zh":
|
|||
|
|
punctuations = ["，", "。", "！", "？"]
|
|||
|
|
elif pred.language == "ja":
|
|||
|
|
punctuations = ["、", "。", "！", "？"]
|
|||
|
|
elif pred.language in ("ar", "fa"):
|
|||
|
|
punctuations = ["،", ".", "!", "؟"]
|
|||
|
|
elif pred.language == "el":
|
|||
|
|
punctuations = [",", ".", "！", "；"]
|
|||
|
|
elif pred.language == "ti":
|
|||
|
|
punctuations = ["།"]
|
|||
|
|
|
|||
|
|
for word in pred.words:
|
|||
|
|
if word.text in punctuations:
|
|||
|
|
punctuation_words.append(word)
|
|||
|
|
start_times = list(map(lambda x: x.start_time, punctuation_words))
|
|||
|
|
start_times = sorted(start_times)
|
|||
|
|
end_times = list(map(lambda x: x.end_time, punctuation_words))
|
|||
|
|
end_times = sorted(end_times)
|
|||
|
|
|
|||
|
|
self.context.punctuation_num = len(self.context.labels)
|
|||
|
|
label_n = len(self.context.labels)
|
|||
|
|
for i, label in enumerate(self.context.labels):
|
|||
|
|
label_left = (label.end - 0.7)
|
|||
|
|
label_right = (label.end + 0.7)
|
|||
|
|
if i < label_n - 1:
|
|||
|
|
label_left = label.end
|
|||
|
|
label_right = self.context.labels[i+1].start
|
|||
|
|
|
|||
|
|
exist = False
|
|||
|
|
|
|||
|
|
def upper_bound(x: float, lst: List[float]) -> int:
|
|||
|
|
ans = -1
|
|||
|
|
left, right = 0, len(lst) - 1
|
|||
|
|
while left <= right:
|
|||
|
|
mid = (left + right) // 2
|
|||
|
|
if lst[mid] >= x:
|
|||
|
|
ans = mid
|
|||
|
|
right = mid - 1
|
|||
|
|
else:
|
|||
|
|
left = mid + 1
|
|||
|
|
return ans
|
|||
|
|
|
|||
|
|
def lower_bound(x: float, lst: List[float]) -> int:
|
|||
|
|
ans = -1
|
|||
|
|
left, right = 0, len(lst) - 1
|
|||
|
|
while left <= right:
|
|||
|
|
mid = (left + right) // 2
|
|||
|
|
if lst[mid] <= x:
|
|||
|
|
ans = mid
|
|||
|
|
left = mid + 1
|
|||
|
|
else:
|
|||
|
|
right = mid - 1
|
|||
|
|
return ans
|
|||
|
|
|
|||
|
|
left_in_pred = upper_bound(label_left, start_times)
|
|||
|
|
if left_in_pred != -1 and start_times[left_in_pred] <= label_right:
|
|||
|
|
exist = True
|
|||
|
|
right_in_pred = lower_bound(label_right, end_times)
|
|||
|
|
if right_in_pred != -1 and end_times[right_in_pred] >= label_left:
|
|||
|
|
exist = True
|
|||
|
|
|
|||
|
|
if exist:
|
|||
|
|
self.context.pred_punctuation_num += 1
|
|||
|
|
return self.context
|
|||
|
|
|
|||
|
|
def _gen_init_data(self) -> dict:
|
|||
|
|
return {
|
|||
|
|
"parameter": {
|
|||
|
|
"lang": None,
|
|||
|
|
"sample_rate": self.context.sample_rate,
|
|||
|
|
"channel": self.context.channel,
|
|||
|
|
"format": self.context.audio_format,
|
|||
|
|
"bits": self.context.bits,
|
|||
|
|
"enable_words": self.context.enable_words,
|
|||
|
|
}
|
|||
|
|
}
|