forked from ufal/whisper_streaming
-
Notifications
You must be signed in to change notification settings - Fork 1
/
ws_client.py
executable file
·126 lines (103 loc) · 3.69 KB
/
ws_client.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
#!/usr/bin/env python
import asyncio
import json
import logging
import signal
import subprocess
import websockets
from websockets import WebSocketClientProtocol
from ws_shared import URL, TranscriptionData
LOG_LEVEL = logging.INFO
logging.basicConfig(
level=LOG_LEVEL, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)
logger.setLevel(LOG_LEVEL)
async def producer(ws: WebSocketClientProtocol):
# Set up the arecord command
arecord_command = [
"arecord",
"-f",
"S16_LE",
"-c1",
"-r",
"16000",
"-t",
"raw",
"-D",
"default",
]
# Open a subprocess that runs the arecord command
process = subprocess.Popen(arecord_command, stdout=subprocess.PIPE)
if process.stdout is None:
raise Exception("stdout is None")
try:
while True:
data = process.stdout.read(8192)
if not data:
break
# Send the audio data through the WebSocket
await ws.send(data)
await asyncio.sleep(0.1)
except Exception as e:
if not isinstance(e, asyncio.CancelledError):
logger.error(
"Received unexpected error", exc_info=True
) # NOTE: will `exc_info` print the stack trace?
finally:
logger.info("Stopping the producer")
process.kill()
process.wait()
logger.info("Killed the arecord process")
def copy_to_clipboard(text: str) -> None:
process = subprocess.Popen(
["xclip", "-selection", "clipboard"], stdin=subprocess.PIPE
)
process.communicate(input=text.encode("utf-8"))
logger.info("Transcript copied to clipboard")
# NOTE: should the process be explicitly killed here?
# TODO: handle the case when the process fails
def send_desktop_notification(title: str, description: str | None = None) -> None:
if description is None:
cmd = ["notify-desktop", title]
else:
cmd = ["notify-desktop", title, description]
_ = subprocess.Popen(cmd)
# NOTE: should the process be explicitly killed here?
logger.info("Sent desktop notification")
def handle_transcription(transcription_data: TranscriptionData) -> None:
if len(transcription_data.transcription) == 0:
return
copy_to_clipboard(transcription_data.transcription)
if transcription_data.is_complete:
send_desktop_notification(transcription_data.transcription)
async def consumer(ws: WebSocketClientProtocol):
async for message in ws:
logger.info(message)
transcription_data = TranscriptionData(**json.loads(message))
handle_transcription(transcription_data)
async def main():
async with websockets.connect(URL) as ws:
loop = asyncio.get_running_loop()
async with asyncio.TaskGroup() as tg:
producer_task = tg.create_task(producer(ws))
tg.create_task(consumer(ws))
async def on_interrupt():
logger.info("Received SIGINT")
producer_task.cancel()
await producer_task # NOTE: doing this for some reason makes all the code below unreacheble
# HACK: using sleep instead
await asyncio.sleep(0.5)
logger.info("Producer stopped")
await ws.send("stop")
logger.info("Sent stop command")
# overrides the default behavior of the SIGINT signal
loop.add_signal_handler(
signal.SIGINT,
lambda: asyncio.create_task(on_interrupt()),
)
if __name__ == "__main__":
try:
asyncio.run(main())
except KeyboardInterrupt:
pass