diff --git a/scripts/asr-streaming-query.py b/scripts/asr-streaming-query.py index 780abd2..4619cbb 100644 --- a/scripts/asr-streaming-query.py +++ b/scripts/asr-streaming-query.py @@ -21,7 +21,6 @@ import websockets # Desired audio properties TARGET_SAMPLE_RATE = 24000 TARGET_CHANNELS = 1 # Mono -HEADERS = {"kyutai-api-key": "open_token"} all_text = [] transcript = [] finished = False @@ -101,10 +100,11 @@ async def send_messages(websocket, rtf: float): print("Connection closed while sending messages.") -async def stream_audio(url: str, rtf: float): +async def stream_audio(url: str, rtf: float, api_key: str): """Stream audio data to a WebSocket server.""" - async with websockets.connect(url, additional_headers=HEADERS) as websocket: + headers = {"kyutai-api-key": api_key} + async with websockets.connect(url, additional_headers=headers) as websocket: send_task = asyncio.create_task(send_messages(websocket, rtf)) receive_task = asyncio.create_task(receive_messages(websocket)) await asyncio.gather(send_task, receive_task) @@ -115,6 +115,7 @@ if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("in_file") parser.add_argument("--transcript") + parser.add_argument("--api-key", default="open_key") parser.add_argument( "--url", help="The url of the server to which to send the audio", @@ -124,7 +125,7 @@ if __name__ == "__main__": args = parser.parse_args() url = f"{args.url}/api/asr-streaming" - asyncio.run(stream_audio(url, args.rtf)) + asyncio.run(stream_audio(url, args.rtf, args.api_key)) print(" ".join(all_text)) if args.transcript is not None: with open(args.transcript, "w") as fobj: