Transcribe speech

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

"""Stream a 16 kHz mono PCM16 WAV to Qwen3-ASR on Baseten.

Requires Python 3.11+ and websockets==15.0.1.
Usage: uv run --python 3.11 transcribe.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

SAMPLE_RATE = 16_000
CHUNK_FRAMES = 1_600  # 100 ms of audio.


def load_audio(path):
    with wave.open(path, "rb") as audio:
        if (audio.getframerate(), audio.getnchannels(), audio.getsampwidth()) != (
            SAMPLE_RATE, 1, 2
        ) or audio.getcomptype() != "NONE":
            raise ValueError("Use an uncompressed 16 kHz mono PCM16 WAV file.")
        pcm = audio.readframes(audio.getnframes())
    if not pcm:
        raise ValueError("The WAV file contains no audio.")
    return pcm


async def transcribe(path):
    pcm = load_audio(path)
    key = os.environ["BASETEN_API_KEY"]
    model_id = os.environ["BASETEN_MODEL_ID"]
    url = f"wss://model-{model_id}.api.baseten.co/environments/production/websocket"
    headers = {"Authorization": f"Bearer {key}"}
    finals = {}

    async with connect(url, additional_headers=headers, open_timeout=60) as ws:
        await ws.send(json.dumps({
            "whisper_params": {"audio_language": "auto"},
            "streaming_params": {
                "enable_partial_transcripts": True,
                "partial_transcript_interval_s": 0.5,
            },
        }))

        async def send_audio():
            start = asyncio.get_running_loop().time()
            for offset in range(0, len(pcm), CHUNK_FRAMES * 2):
                chunk = pcm[offset:offset + CHUNK_FRAMES * 2]
                await ws.send(json.dumps({
                    "type": "input_audio_buffer.append",
                    "audio": base64.b64encode(chunk).decode("ascii"),
                }))
                # Pace by audio duration without accumulating network delays.
                next_send = start + (offset + len(chunk)) / (SAMPLE_RATE * 2)
                await asyncio.sleep(max(0, next_send - asyncio.get_running_loop().time()))
            await ws.send(json.dumps({"type": "input_audio_buffer.commit"}))

        async def receive_transcripts():
            while True:
                message = json.loads(await asyncio.wait_for(ws.recv(), timeout=60))
                if message.get("type") == "error" or message.get("error"):
                    raise RuntimeError(f"Server error: {message}")
                if message.get("type") != "transcription":
                    continue
                number = message["transcription_num"]
                text = " ".join(segment["text"] for segment in message.get("segments", []))
                if message.get("is_final"):
                    finals[number] = text
                    print(f"final {number}: {text}", flush=True)
                    if message.get("is_end_of_audio_flush"):
                        return
                else:
                    # Each update replaces the hypothesis for this number.
                    print(f"partial {number} (revisable): {text}", flush=True)

        sender = asyncio.create_task(send_audio())
        receiver = asyncio.create_task(receive_transcripts())
        try:
            # Bound the whole session even if unrelated events keep arriving.
            async with asyncio.timeout(len(pcm) / (SAMPLE_RATE * 2) + 120):
                await asyncio.gather(sender, receiver)
        finally:
            for task in (sender, receiver):
                task.cancel()
            await asyncio.gather(sender, receiver, return_exceptions=True)

    return " ".join(finals[number] for number in sorted(finals) if finals[number])


if __name__ == "__main__":
    if sys.version_info < (3, 11):
        sys.exit("This client requires Python 3.11 or later.")
    if len(sys.argv) != 2:
        sys.exit("Usage: uv run --python 3.11 transcribe.py <16-kHz-mono-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. Check deployment readiness and model logs.")
    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.")