NVIDIA Nemotron 3.5 ASR Streaming 0.6B (English)

1import asyncio, base64, json, wave, websockets
2
3URL = "wss://model-<id>.api.baseten.co/environments/production/websocket?intent=transcription"
4API_KEY = "<baseten-api-key>"
5MODEL = "cache-aware-parakeet-rnnt-en-US-asr-streaming-sortformer"   # served ASR model name
6SAMPLE_RATE = 16000
7CHUNK_MS = 250
8
9def load_pcm(path):
10    # Return mono 16-bit PCM at SAMPLE_RATE. The server treats the bytes as raw
11    # pcm16 at the declared rate, so any stereo / non-16kHz input must be
12    # converted here or it transcribes to nothing.
13    with wave.open(path, "rb") as wf:
14        n_ch, width, rate = wf.getnchannels(), wf.getsampwidth(), wf.getframerate()
15        frames = wf.readframes(wf.getnframes())
16    if width != 2:
17        raise ValueError(f"need 16-bit PCM, got {width * 8}-bit")
18    if n_ch == 2:
19        frames = audioop.tomono(frames, 2, 0.5, 0.5)
20    if rate != SAMPLE_RATE:
21        frames, _ = audioop.ratecv(frames, 2, 1, rate, SAMPLE_RATE, None)
22    return frames
23
24async def main(path):
25    pcm = load_pcm(path)
26    chunk = int(SAMPLE_RATE * CHUNK_MS / 1000) * 2   # 2 bytes/sample
27
28    async with websockets.connect(URL, additional_headers={"Authorization": f"Api-Key {API_KEY}"}) as ws:
29        # 1. Configure the session (model, language, audio format, sample rate).
30        await ws.send(json.dumps({
31            "type": "transcription_session.update",
32            "session": {
33                "input_audio_format": "pcm16",
34                "input_audio_transcription": {"model": MODEL, "language": "en-US"},
35                "input_audio_params": {"sample_rate_hz": SAMPLE_RATE, "num_channels": 1},
36                "recognition_config": {
37                    "max_alternatives": 1,
38                    "enable_automatic_punctuation": True,
39                },
40            },
41        }))
42
43        async def send():
44            # 2. Stream base64 PCM chunks; commit each one to get streaming output.
45            for i in range(0, len(pcm), chunk):
46                b64 = base64.b64encode(pcm[i:i + chunk]).decode()
47                await ws.send(json.dumps({"type": "input_audio_buffer.append", "audio": b64}))
48                await ws.send(json.dumps({"type": "input_audio_buffer.commit"}))
49                await asyncio.sleep(CHUNK_MS / 1000)     # pace like real-time mic
50            # 3. Flush + stop (required when sending a finite file).
51            await ws.send(json.dumps({"type": "input_audio_buffer.done"}))
52
53        async def receive():
54            # 4. delta = cumulative interim partial; completed = finalized segment.
55            async for raw in ws:
56                msg = json.loads(raw)
57                t = msg.get("type")
58                if t == "conversation.item.input_audio_transcription.delta":
59                    print("partial:", msg.get("delta", ""))
60                elif t == "conversation.item.input_audio_transcription.completed":
61                    print("final:  ", msg.get("transcript", ""))
62                    if msg.get("is_last_result"):
63                        break
64
65        await asyncio.gather(send(), receive())
66
67asyncio.run(main("audio.wav"))