-
Notifications
You must be signed in to change notification settings - Fork 3
/
orca_demo_streaming.py
335 lines (280 loc) · 11 KB
/
orca_demo_streaming.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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
#
# Copyright 2024 Picovoice Inc.
#
# You may not use this file except in compliance with the license. A copy of the license is located in the "LICENSE"
# file accompanying this source.
#
# Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
# an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
# specific language governing permissions and limitations under the License.
#
import argparse
import platform
import re
import subprocess
import threading
import time
from dataclasses import dataclass
from queue import Queue
from collections import deque
from itertools import chain
from typing import (
Callable,
Optional,
Sequence,
)
import pvorca
import tiktoken
from pvorca import (
Orca,
OrcaActivationLimitError,
OrcaInvalidArgumentError,
)
from pvspeaker import PvSpeaker
CUSTOM_PRON_PATTERN = r"\{(.*?\|.*?)\}"
CUSTOM_PRON_PATTERN_NO_WHITESPACE = r"\{(.*?\|.*?)\}(?!\s)"
def linux_machine() -> str:
machine = platform.machine()
if machine == "x86_64":
return machine
elif machine in ["aarch64", "armv7l"]:
arch_info = ("-" + machine) if "64bit" in platform.architecture()[0] else ""
else:
raise NotImplementedError("Unsupported CPU architecture: `%s`" % machine)
cpu_info = ""
try:
cpu_info = subprocess.check_output(["cat", "/proc/cpuinfo"]).decode("utf-8")
cpu_part_list = [x for x in cpu_info.split("\n") if "CPU part" in x]
cpu_part = cpu_part_list[0].split(" ")[-1].lower()
except Exception as e:
raise RuntimeError("Failed to identify the CPU with `%s`\nCPU info: `%s`" % (e, cpu_info))
if "0xd03" == cpu_part:
return "cortex-a53" + arch_info
elif "0xd08" == cpu_part:
return "cortex-a72" + arch_info
elif "0xd0b" == cpu_part:
return "cortex-a76" + arch_info
else:
raise NotImplementedError("Unsupported CPU: `%s`." % cpu_part)
class OrcaThread:
@dataclass
class OrcaInput:
text: str
flush: bool
def __init__(
self,
orca: Orca,
flush_audio_callback: Callable[[Sequence[int]], None],
play_audio_callback: Callable[[Sequence[int]], int],
num_tokens_per_second: int,
audio_wait_chunks: Optional[int] = None,
) -> None:
self._orca = orca
self._orca_stream = self._orca.stream_open()
self._play_audio_callback = play_audio_callback
self._flush_audio_callback = flush_audio_callback
self._num_tokens_per_second = num_tokens_per_second
assert self._num_tokens_per_second > 0
self._queue: Queue[Optional[OrcaThread.OrcaInput]] = Queue()
self._thread = None
self._time_first_audio_available = -1
self._pcm_buffer = deque()
self._wait_chunks = audio_wait_chunks or self._get_first_audio_wait_chunks()
self._num_pcm_chunks_processed = 0
@staticmethod
def _get_first_audio_wait_chunks() -> int:
wait_chunks = 0
if platform.system() == "Linux":
machine = linux_machine()
if "cortex" in machine:
wait_chunks = 1
return wait_chunks
def _run(self) -> None:
while True:
orca_input = self._queue.get()
if orca_input is None:
break
try:
if not orca_input.flush:
pcm = self._orca_stream.synthesize(orca_input.text)
else:
pcm = self._orca_stream.flush()
except OrcaInvalidArgumentError as e:
raise ValueError(f"Orca could not synthesize text input `{orca_input.text}`: `{e}`")
if pcm is not None:
if self._num_pcm_chunks_processed == 0:
self._time_first_audio_available = time.time()
self._num_pcm_chunks_processed += 1
self._pcm_buffer.append(pcm)
if self._num_pcm_chunks_processed > self._wait_chunks:
if len(self._pcm_buffer) > 0:
pcm = self._pcm_buffer.popleft()
written = self._play_audio_callback(pcm)
if written < len(pcm):
self._pcm_buffer.appendleft(pcm[written:])
def _close_thread_blocking(self):
self._queue.put_nowait(None)
self._thread.join()
def start(self) -> None:
self._thread = threading.Thread(target=self._run)
self._thread.start()
def synthesize(self, text: str) -> None:
self._queue.put_nowait(self.OrcaInput(text=text, flush=False))
def flush(self) -> None:
self._queue.put_nowait(self.OrcaInput(text="", flush=True))
self._close_thread_blocking()
def flush_audio(self) -> None:
remaining_pcm = list(chain.from_iterable(self._pcm_buffer))
self._thread = threading.Thread(target=self._flush_audio_callback, args=(remaining_pcm,))
self._thread.start()
self._thread.join()
def delete(self) -> None:
self._close_thread_blocking()
self._orca_stream.close()
self._orca.delete()
def get_time_first_audio_available(self) -> float:
return self._time_first_audio_available
def tokenize_text(text: str) -> Sequence[str]:
text = re.sub(CUSTOM_PRON_PATTERN_NO_WHITESPACE, r'{\1} ', text)
custom_pronunciations = re.findall(CUSTOM_PRON_PATTERN, text)
custom_pronunciations = set(["{" + pron + "}" for pron in custom_pronunciations])
encoder = tiktoken.encoding_for_model("gpt-4")
tokens_raw = [encoder.decode([i]) for i in encoder.encode(text)]
custom_pron = ""
tokens_with_custom_pronunciations = []
for i, token in enumerate(tokens_raw):
in_custom_pron = False
for pron in custom_pronunciations:
in_custom_pron_global = len(custom_pron) > 0
current_match = token.strip() if not in_custom_pron_global else custom_pron + token
if pron.startswith(current_match):
custom_pron += token.strip() if not in_custom_pron_global else token
in_custom_pron = True
if not in_custom_pron:
if custom_pron != "":
tokens_with_custom_pronunciations.append(f" {custom_pron}" if i != 0 else custom_pron)
custom_pron = ""
tokens_with_custom_pronunciations.append(token)
return tokens_with_custom_pronunciations
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument(
"--access_key",
"-a",
required=True,
help="AccessKey obtained from Picovoice Console (https://console.picovoice.ai/)")
parser.add_argument(
"--library_path",
"-l",
help="Absolute path to dynamic library. Default: using the library provided by `pvorca`")
parser.add_argument(
"--model_path",
"-m",
help="Absolute path to Orca model. Default: using the model provided by `pvorca`")
parser.add_argument(
"--text_to_stream",
"-t",
required=True,
help="Text to be streamed to Orca")
parser.add_argument(
"--tokens_per_second",
type=int,
default=15,
help="Number of tokens per second to be streamed to Orca, simulating an LLM response.")
parser.add_argument(
"--audio_wait_chunks",
type=int,
default=None,
help="Number of PCM chunks to wait before starting to play audio. Default: system-dependent.")
parser.add_argument(
"--buffer_size_secs",
type=int,
default=20,
help="The size in seconds of the internal buffer used by pvspeaker to play audio.")
parser.add_argument(
"--show_audio_devices",
action="store_true",
help="Only list available audio output devices and exit")
parser.add_argument(
'--audio-device-index',
type=int,
default=-1,
help='Index of input audio device')
args = parser.parse_args()
if args.show_audio_devices:
devices = PvSpeaker.get_available_devices()
for i in range(len(devices)):
print("index: %d, device name: %s" % (i, devices[i]))
exit(0)
access_key = args.access_key
model_path = args.model_path
library_path = args.library_path
text = args.text_to_stream
tokens_per_second = args.tokens_per_second
audio_wait_chunks = args.audio_wait_chunks
buffer_size_secs = args.buffer_size_secs
audio_device_index = args.audio_device_index
orca = pvorca.create(access_key=access_key, model_path=model_path, library_path=library_path)
speaker = None
try:
speaker = PvSpeaker(
sample_rate=orca.sample_rate,
bits_per_sample=16,
buffer_size_secs=buffer_size_secs,
device_index=audio_device_index)
speaker.start()
except RuntimeError or ValueError:
print(
"\nWarning: Failed to initialize PvSpeaker. Orca will still generate PCM data, "
"but it will not be played.\n")
def play_audio_callback(pcm: Sequence[int]) -> int:
try:
if speaker is not None:
return speaker.write(pcm)
return len(pcm)
except ValueError:
pass
return len(pcm)
def flush_audio_callback(pcm: Sequence[int]) -> None:
try:
if speaker is not None:
speaker.flush(pcm)
except MemoryError:
pass
orca_thread = OrcaThread(
orca=orca,
play_audio_callback=play_audio_callback,
flush_audio_callback=flush_audio_callback,
num_tokens_per_second=tokens_per_second,
audio_wait_chunks=audio_wait_chunks,
)
orca_thread.start()
try:
print(f"Orca version: {orca.version}\n")
tokens = tokenize_text(text=text)
print(f"Simulated text stream:")
time_start_text_stream = time.time()
for token in tokens:
print(f"{token}", end="", flush=True)
orca_thread.synthesize(text=token)
time.sleep(1 / tokens_per_second)
text_stream_duration_seconds = time.time() - time_start_text_stream
orca_thread.flush()
first_audio_available_seconds = orca_thread.get_time_first_audio_available() - time_start_text_stream
print(f"\n\nTime to finish text stream: {text_stream_duration_seconds:.2f} seconds")
print(f"Time to receive first audio: {first_audio_available_seconds:.2f} seconds after text stream started\n")
if speaker is not None:
print("Waiting for audio to finish ...")
orca_thread.flush_audio()
if speaker is not None:
speaker.delete()
except KeyboardInterrupt:
print("\nStopped...")
if speaker is not None:
speaker.stop()
except OrcaActivationLimitError:
print("\nAccessKey has reached its processing limit")
finally:
orca_thread.delete()
if __name__ == "__main__":
main()