Transcribe streaming audio

import asyncio
import signal
import websockets
import sounddevice as sd
import numpy as np
import json
import os

SAMPLE_RATE = 16000
CHUNK_SIZE = 512  # ~32ms per chunk

headers = {"Authorization": f"Bearer {os.environ['BASETEN_API_KEY']}"}
ws_url = "wss://model-{model_id}.api.baseten.co/environments/production/websocket"
metadata = {
    "streaming_params": {"encoding": "pcm_s16le", "sample_rate": SAMPLE_RATE},
    "whisper_params": {"audio_language": "en"},
}

async def stream_mic():
    loop = asyncio.get_running_loop()
    send_queue = asyncio.Queue()
    stop_event = asyncio.Event()

    def audio_callback(indata, frames, time_info, status):
        loop.call_soon_threadsafe(
            send_queue.put_nowait, (indata * 32767).astype(np.int16).tobytes()
        )

    # Ctrl+C sets stop_event instead of raising KeyboardInterrupt,
    # so the end_audio handshake can complete cleanly before closing.
    loop.add_signal_handler(signal.SIGINT, lambda: loop.call_soon_threadsafe(stop_event.set))

    async with websockets.connect(ws_url, additional_headers=headers) as ws:
        await ws.send(json.dumps(metadata))
        print("Recording. Press Ctrl+C to stop.\n")

        async def send_audio():
            with sd.InputStream(samplerate=SAMPLE_RATE, blocksize=CHUNK_SIZE,
                                channels=1, dtype="float32", callback=audio_callback):
                while not stop_event.is_set():
                    try:
                        chunk = await asyncio.wait_for(send_queue.get(), timeout=0.1)
                        await ws.send(chunk)
                    except asyncio.TimeoutError:
                        continue
            # Drain any chunks buffered after stop
            while not send_queue.empty():
                await ws.send(send_queue.get_nowait())
            # Flush remaining speech buffered on the server
            await ws.send(json.dumps({"type": "end_audio"}))

        async def receive_messages():
            # Receive concurrently with send_audio; VAD may trigger transcription
            # results while audio is still being sent; sequential receive would miss them.
            async for raw in ws:
                msg = json.loads(raw)
                if msg.get("type") == "transcription":
                    text = " ".join(s["text"] for s in msg.get("segments", []))
                    print(f"[{'final' if msg['is_final'] else 'partial'}] {text}")
                elif msg.get("type") == "end_audio":
                    if msg.get("body", {}).get("status") == "finished":
                        break  # All audio processed; safe to close

        await asyncio.gather(send_audio(), receive_messages())

asyncio.run(stream_mic())

import asyncio
import time
import websockets
import json
import os

SAMPLE_RATE = 16000
CHUNK_SAMPLES = 512          # ~32ms per chunk, matches live mic cadence
CHUNK_SIZE = CHUNK_SAMPLES * 2  # bytes (pcm_s16le = 2 bytes/sample)
CHUNK_DURATION_S = CHUNK_SAMPLES / SAMPLE_RATE

headers = {"Authorization": f"Bearer {os.environ['BASETEN_API_KEY']}"}
ws_url = "wss://model-{model_id}.api.baseten.co/environments/production/websocket"


async def transcribe_session(audio_bytes: bytes, language: str = "en") -> str:
    """Open a new connection, transcribe one audio buffer, close cleanly."""
    metadata = {
        "streaming_params": {"encoding": "pcm_s16le", "sample_rate": SAMPLE_RATE},
        "whisper_params": {"audio_language": language},
    }
    transcripts = []

    async with websockets.connect(ws_url, additional_headers=headers) as ws:
        await ws.send(json.dumps(metadata))

        async def send_audio():
            for i in range(0, len(audio_bytes), CHUNK_SIZE):
                chunk = audio_bytes[i : i + CHUNK_SIZE]
                # VAD requires ≥ 512 samples per chunk. Zero-pad the last
                # chunk if the file doesn't divide evenly.
                if len(chunk) < CHUNK_SIZE:
                    chunk = chunk + b"\x00" * (CHUNK_SIZE - len(chunk))
                await ws.send(chunk)
                # Pace at real-time speed so VAD sees audio at the same
                # cadence as a live mic; sending faster causes idle timeouts.
                await asyncio.sleep(CHUNK_DURATION_S)
            await ws.send(json.dumps({"type": "end_audio"}))

        async def receive_messages():
            # Receive concurrently; VAD may emit transcripts while audio
            # is still being sent; sequential receive would miss them.
            async for raw in ws:
                message = json.loads(raw)
                if message.get("type") == "transcription":
                    transcripts.append(
                        " ".join(s["text"] for s in message.get("segments", []))
                    )
                elif message.get("type") == "end_audio":
                    if message.get("body", {}).get("status") == "finished":
                        break

        await asyncio.gather(send_audio(), receive_messages())

    return " ".join(transcripts)


async def process_sequential(audio_files: list[bytes]):
    """One connection per file, each opened after the previous one closes."""
    for audio in audio_files:
        transcript = await transcribe_session(audio)
        print(f"Transcript: {transcript}")


async def process_overlapping(audio_files: list[bytes]):
    """All connections opened in parallel; wall-clock time ≈ longest file."""
    results = await asyncio.gather(*[transcribe_session(a) for a in audio_files])
    for transcript in results:
        print(f"Transcript: {transcript}")