Whisper streaming

1
import websockets
2
import asyncio
3
import json
4
from urllib.parse import urlencode
5
6
import pyaudio
7
8
FRAMES_PER_BUFFER = 3200
9
FORMAT = pyaudio.paInt16
10
CHANNELS = 1
11
RATE = 48000
12
p = pyaudio.PyAudio()
13
14
stream = p.open(
15
    format=FORMAT,
16
    channels=CHANNELS,
17
    rate=RATE,
18
    input=True,
19
    frames_per_buffer=FRAMES_PER_BUFFER
20
)
21
22
BASE_URL = "wss://streaming.assemblyai.com/v3/ws"
23
CONNECTION_PARAMS = {
24
    "sample_rate": RATE,
25
    "speech_model": "whisper-rt",
26
    "language_detection": True,
27
}
28
URL = f"{BASE_URL}?{urlencode(CONNECTION_PARAMS)}"
29
30
async def send_receive():
31
32
    print(f'Connecting websocket to url ${URL}')
33
34
    async with websockets.connect(
35
        URL,
36
        additional_headers={"Authorization": "YOUR-API-KEY"},
37
        ping_interval=5,
38
        ping_timeout=20
39
    ) as _ws:
40
        await asyncio.sleep(0.1)
41
        print("Receiving SessionBegins ...")
42
43
        session_begins = await _ws.recv()
44
        print(session_begins)
45
        print("Sending messages ...")
46
47
        async def send():
48
            while True:
49
                try:
50
                    data = stream.read(FRAMES_PER_BUFFER, exception_on_overflow=False)
51
                    await _ws.send(data)
52
                except websockets.exceptions.ConnectionClosedError as e:
53
                    print(e)
54
                except Exception as e:
55
                    print(e)
56
                await asyncio.sleep(0.01)
57
58
        async def receive():
59
            while True:
60
                try:
61
                    result_str = await _ws.recv()
62
                    data = json.loads(result_str)
63
                    transcript = data['transcript']
64
                    utterance = data['utterance']
65
66
                    if data['type'] == 'Turn':
67
                        if not data.get('end_of_turn') and transcript:
68
                            print(f"[PARTIAL TURN TRANSCRIPT]: {transcript}")
69
                        if data.get('utterance'):
70
                            print(f"[PARTIAL TURN UTTERANCE]: {utterance}")
71
                            # Display language detection info if available
72
                            if 'language_code' in data:
73
                                print(f"[UTTERANCE LANGUAGE DETECTION]: {data['language_code']} - {data['language_confidence']:.2%}")
74
                        if data.get('end_of_turn'):
75
                            print(f"[FULL TURN TRANSCRIPT]: {transcript}")
76
                            # Display language detection info if available
77
                            if 'language_code' in data:
78
                                print(f"[END OF TURN LANGUAGE DETECTION]: {data['language_code']} - {data['language_confidence']:.2%}")
79
                    else:
80
                        pass
81
82
                except websockets.exceptions.ConnectionClosed:
83
                    break
84
                except Exception as e:
85
                    print(f"\nError receiving data: {e}")
86
                    break
87
88
        try:
89
            await asyncio.gather(send(), receive())
90
        except KeyboardInterrupt:
91
            await _ws.send({"type": "Terminate"})
92
            # Wait for the server to close the connection after receiving the message
93
            await _ws.wait_closed()
94
            print("Session terminated and connection closed.")
95
96
if __name__ == "__main__":
97
    try:
98
        asyncio.run(send_receive())
99
    finally:
100
        stream.stop_stream()
101
        stream.close()
102
        p.terminate()