whispering/whisper_streaming/serve.py
Yuta Hayashibe c6f3e1e86b Fix for lint
2022-09-24 21:52:27 +09:00

53 lines
1.2 KiB
Python

#!/usr/bin/env python3
import asyncio
from logging import getLogger
import numpy as np
import websockets
from whisper_streaming.transcriber import WhisperStreamingTranscriber
logger = getLogger(__name__)
async def serve_with_websocket_main(websocket):
global g_wsp
idx: int = 0
while True:
logger.debug(f"Segment #: {idx}")
message = await websocket.recv()
if isinstance(message, str):
logger.debug(f"Got str: {message}")
continue
logger.debug(f"Message size: {len(message)}")
segment = np.frombuffer(message, dtype=np.float32)
for chunk in g_wsp.transcribe(segment=segment):
await websocket.send(chunk.json())
idx += 1
async def serve_with_websocket(
*,
wsp: WhisperStreamingTranscriber,
host: str,
port: int,
):
logger.info(f"Serve at {host}:{port}")
logger.info("Make secure with your responsibility!")
global g_wsp
g_wsp = wsp
try:
async with websockets.serve( # type: ignore
serve_with_websocket_main,
host=host,
port=port,
max_size=999999999,
):
await asyncio.Future()
except KeyboardInterrupt:
pass