Streaming Transcription

!pip3 install requests torch torchaudio websocket-client

import io
import time
import json
import torch
import requests
import torchaudio
import threading
import websocket
import urllib.parse

lock = threading.Lock()
state = {}

def on_open(ws):
    def send_audio_chunks():
        for chunk in audio_chunk_bytes:
            ws.send(chunk, opcode=websocket.ABNF.OPCODE_BINARY)
            time.sleep(chunk_size_ms / 1000)

        final_checkpoint = json.dumps({"checkpoint_id": "final"})
        ws.send(final_checkpoint, opcode=websocket.ABNF.OPCODE_TEXT)

    threading.Thread(target=send_audio_chunks).start()

def on_message(ws, message):
    message = json.loads(message)
    if message.get("checkpoint_id") == "final":
        ws.close()
        return

    update = {s["id"]: s["text"] for s in message["segments"]}
    with lock:
        state.update(update)
        print("\n".join(f" - {k}: {v}" for k, v in state.items()))

def on_error(ws, error):
    print(f"WebSocket error: {error}")

# Open a connection URL with query params
url = "ws://audio-streaming.us-virginia-1.direct.fireworks.ai/v1/audio/transcriptions/streaming"
params = urllib.parse.urlencode({
    "language": "en",
})
ws = websocket.WebSocketApp(
    f"{url}?{params}",
    header={"Authorization": "<FIREWORKS_API_KEY>"},
    on_open=on_open,
    on_message=on_message,
    on_error=on_error,
)
ws.run_forever()