Transcribe speech

# /// script
# requires-python = ">=3.11"
# dependencies = ["websockets"]
# ///

"""Stream a PCM16 WAV to a unified Baseten STT deployment.

Usage: uv run --python 3.11 transcribe-v2.py jfk.wav
"""

import asyncio
import base64
import json
import os
import sys
import wave

from websockets.asyncio.client import connect
from websockets.exceptions import WebSocketException


async def receive_event(ws, expected=None):
    event = json.loads(await asyncio.wait_for(ws.recv(), timeout=60))
    if event["type"] == "error":
        raise RuntimeError(f"Server error: {event['code']}: {event['message']}")
    if expected is not None and event["type"] != expected:
        raise RuntimeError(f"Expected {expected}, received {event['type']}")
    return event


async def transcribe(path):
    with wave.open(path, "rb") as audio:
        audio_format = (audio.getframerate(), audio.getnchannels(), audio.getsampwidth())
        if audio.getcomptype() != "NONE":
            raise ValueError("Use an uncompressed PCM16 WAV file.")
        pcm = audio.readframes(audio.getnframes())
    if not pcm:
        raise ValueError("The WAV file contains no audio.")

    key = os.environ["BASETEN_API_KEY"]
    model_id = os.environ["BASETEN_MODEL_ID"]
    deployment_id = os.environ["BASETEN_DEPLOYMENT_ID"]
    url = f"wss://model-{model_id}.api.baseten.co/deployment/{deployment_id}/websocket"
    headers = {"Authorization": f"Bearer {key}"}
    finals = []

    # Allow a scaled-to-zero deployment time to start.
    async with connect(url, additional_headers=headers, open_timeout=600) as ws:
        created = await receive_event(ws, "session.created")
        caps = created["capabilities"]
        fmt = caps["input_format"]
        if audio_format != (fmt["sample_rate_hz"], fmt["channels"], 2):
            raise ValueError(f"Convert the WAV to the deployment's format: {fmt}")
        if "client_flush" not in caps["segmentation_modes"]:
            raise RuntimeError("This client requires client_flush support.")

        settings = {"segmentation_mode": "client_flush"}
        if caps["partials"]:
            settings["partial_interval_s"] = 0.5
        await ws.send(json.dumps({"type": "session.update", "session": settings}))
        updated = await receive_event(ws, "session.updated")
        for refusal in updated["refusals"]:
            print("refused:", refusal, flush=True)
        if updated["granted"]["segmentation_mode"] != "client_flush":
            raise RuntimeError("The deployment didn't grant client_flush.")

        rate = fmt["sample_rate_hz"]
        frame_bytes = fmt["channels"] * 2
        chunk_bytes = max(1, rate // 10) * frame_bytes  # 100 ms of audio.
        # Flush before any automatic max-turn cut, then resolve this turn.
        turn_s = min(10.0, caps.get("max_turn_s", 20.0) / 2)
        turn_frames = int(rate * turn_s)
        if turn_frames < 1:
            raise ValueError("The deployment's turn limit is too short for this client.")
        turn_bytes = turn_frames * frame_bytes

        async def send_turn(audio):
            start = asyncio.get_running_loop().time()
            for offset in range(0, len(audio), chunk_bytes):
                chunk = audio[offset:offset + chunk_bytes]
                await ws.send(json.dumps({
                    "type": "input_audio_buffer.append",
                    "audio": base64.b64encode(chunk).decode("ascii"),
                }))
                next_send = start + (offset + len(chunk)) / (rate * frame_bytes)
                await asyncio.sleep(max(0, next_send - asyncio.get_running_loop().time()))
            await ws.send(json.dumps({"type": "input_audio_buffer.flush"}))

        async def receive_turn(turn_id):
            while True:
                event = await receive_event(ws)
                if event["type"] == "transcription.failed":
                    raise RuntimeError(f"Turn {event['turn_id']} failed: {event['message']}")
                if event["type"] not in ("transcription.delta", "transcription.done"):
                    raise RuntimeError(f"Unexpected event: {event['type']}")
                transcript = event["transcript"]
                if transcript["turn_id"] != turn_id or event.get("ended_by"):
                    raise RuntimeError("Unexpected turn boundary; transcription is incomplete.")
                text = "".join(segment["text"] for segment in transcript["segments"]).strip()
                if event["type"] == "transcription.done":
                    print(f"final {turn_id}: {text}", flush=True)
                    return text
                print(f"partial {turn_id}: {text}", flush=True)

        for turn_id, offset in enumerate(range(0, len(pcm), turn_bytes)):
            audio = pcm[offset:offset + turn_bytes]
            sender = asyncio.create_task(send_turn(audio))
            receiver = asyncio.create_task(receive_turn(turn_id))
            try:
                # Bound the whole turn even if partial events keep arriving.
                async with asyncio.timeout(len(audio) / (rate * frame_bytes) + 120):
                    await asyncio.gather(sender, receiver)
            finally:
                for task in (sender, receiver):
                    task.cancel()
                await asyncio.gather(sender, receiver, return_exceptions=True)
            finals.append(receiver.result())

        # Every flushed turn has a final result before close starts the drain.
        await ws.send(json.dumps({"type": "close"}))
        await asyncio.wait_for(ws.wait_closed(), timeout=10)
        if ws.close_code != 1000:
            raise RuntimeError(f"Unexpected WebSocket close code: {ws.close_code}")

    return " ".join(text for text in finals if text)


if __name__ == "__main__":
    if len(sys.argv) != 2:
        sys.exit("Usage: uv run --python 3.11 transcribe-v2.py <PCM16.wav>")
    try:
        print("Transcript:", asyncio.run(transcribe(sys.argv[1])))
    except KeyError as error:
        sys.exit(f"Missing environment variable or response field: {error}")
    except TimeoutError:
        sys.exit("Transcription timed out; a complete transcript wasn't received.")
    except (OSError, ValueError, EOFError, wave.Error, RuntimeError, WebSocketException) as error:
        sys.exit(f"Transcription failed: {error}")
    except KeyboardInterrupt:
        sys.exit("Stopped. Run the file again to receive a complete transcript.")