Add nemotron-3-diarization endpoint docs and Python clients

nemotron-3-diarization/README.md

@@ -0,0 +1,92 @@
+# NVIDIA Nemotron 3 Diarization on Baseten
+
+[NVIDIA Nemotron 3 Diarization](https://huggingface.co/nvidia/Nemotron-3-Diarization) answers
+"who spoke when" in real-world audio: up to eight speakers, labels ordered by first arrival, and a
+single checkpoint that runs at any of several algorithmic latencies, from a 0.32-second input
+buffer for live agents to a 30-second buffer for recordings. It is the successor to Streaming
+Sortformer v2.1 and is released under the [OpenMDW 1.1](https://openmdw.ai/license/1-1/) license,
+which permits commercial use.
+
+Baseten serves the same model three ways. Each preset runs on one NVIDIA RTX PRO 6000 behind a
+serving engine built around NVIDIA's NeMo streaming loop, so throughput and concurrency exceed
+the reference script on the same GPU.
+
+| Preset | Transport | Use it for | Capacity per GPU |
+|---|---|---|---|
+| [`batch/`](batch/) | HTTP, one request per file | Recorded audio: a file URL or base64 in, speaker turns out | 200 six-minute files per minute, sustained |
+| [`streaming/`](streaming/) | WebSocket, one connection per stream | Live audio: 100 ms PCM frames in, a live turn list on every chunk | 560 hour-long streams at the 1.04 s profile, 200 at 0.32 s |
+| [`diarized-transcription-streaming/`](diarized-transcription-streaming/) | WebSocket | Live speaker-tagged words: the diarizer paired with NVIDIA's multitalker Parakeet 0.6B ASR (English) | 190 hour-long streams |
+
+A multilingual transcription pairing is in progress and will be added as a fourth preset.
+
+Model Library page: <https://www.baseten.co/library/nemotron-3-diarization>
+
+## Latency profiles
+
+One checkpoint, chosen per request (batch) or per connection (streaming). Input-buffer latency is
+(chunk + right context) × 80 ms; compute runs far faster than real time, so the buffer dominates
+the time from speech to label.
+
+| Profile | Input-buffer latency | Notes |
+|---|---|---|
+| `offline` | 30.4 s | Best accuracy. Default for batch. |
+| `low` | 1.04 s | Real-time default. |
+| `ultralow` | 0.32 s | Lowest latency; three times the step rate of `low`. |
+
+## Accuracy
+
+Diarization error rate with overlapping speech included and no collar, on identical audio and
+references for every system. Lower is better. Measured on the General Access checkpoint.
+
+| Dataset | Nemotron 3 `offline` | `low` | `ultralow` | Streaming Sortformer v2.1 (`offline`) |
+|---|---|---|---|---|
+| NOTSOFAR-1 (129 meetings) | 15.8 | 17.5 | 18.8 | 23.8 |
+| AMI-SDM (34 meetings) | 25.0 | 25.5 | 25.9 | 27.8 |
+| CALLHOME (12 calls) | 13.0 | 13.8 | 14.0 | 16.7 |
+| AISHELL-4, Mandarin (12 meetings) | 10.2 | 9.8 | 11.8 | 27.2 |
+
+Nemotron 3 at its fastest profile beats its predecessor at its slowest on every set. The gain is
+largest on recordings with five or more speakers, where v2.1's four-speaker cap costs it.
+
+Real-time diarized transcription (English pairing) scores 28.6 cpWER on NOTSOFAR-1 eval-30 with the
+canonical scorer. Words appear a median 0.7 s after they are spoken, speaker labels are final on
+arrival, and committed words are never revised.
+
+## Deploying
+
+Deploy any preset in one click from the Baseten Model Library; each deploy link creates the model in
+your workspace on an NVIDIA RTX PRO 6000:
+
+- Batch diarization — <https://app.baseten.co/deploy/baseten/nemotron-3-diarization-batch>
+- Streaming diarization — <https://app.baseten.co/deploy/baseten/nemotron-3-diarization-streaming>
+- Real-time diarized transcription — <https://app.baseten.co/deploy/baseten/nemotron-3-diarized-transcription-streaming>
+
+The checkpoint (`nvidia/Nemotron-3-Diarization`, OpenMDW 1.1) is fetched at deploy time.
+
+## Calling the endpoints
+
+Each preset directory documents its full request and response contract and ships runnable
+clients (`client.py`, and `curl.sh` for batch). All of them read two environment variables:
+
+```bash
+export BASETEN_API_KEY=...   # https://app.baseten.co/settings/account/api_keys
+export MODEL_ID=...          # from the model's page in the Baseten dashboard
+```
+
+URL forms:
+
+| Preset | Production URL |
+|---|---|
+| batch | `https://model-{MODEL_ID}.api.baseten.co/environments/production/predict` |
+| streaming, diarized transcription | `wss://model-{MODEL_ID}.api.baseten.co/environments/production/websocket` |
+
+Replace `environments/production` with `deployment/{DEPLOYMENT_ID}` to target a specific
+deployment. Authenticate with an `Authorization: Api-Key $BASETEN_API_KEY` header.
+
+Audio for the WebSocket presets is 16 kHz mono PCM16, base64-encoded inside JSON text frames
+(100 ms frames, 3,200 bytes each, work well). The batch preset accepts any streamable format `ffmpeg`
+decodes (WAV, FLAC, MP3, OGG, WebM) and resamples it server-side; see [`batch/README.md`](batch/README.md)
+for the MP4/M4A and URL-fetch caveats.
+
+Every client and response in these directories was run against a live deployment of this code on
+2026-09-22; the JSON shown is what came back, trimmed only for length.

nemotron-3-diarization/batch/README.md

@@ -0,0 +1,139 @@
+# Nemotron 3 Diarization — batch (HTTP)
+
+Diarize a recording in one request: a file URL or base64 audio in, speaker turns out. The latency
+profile is chosen per request. Independent requests arriving within ~100 ms are batched into one
+GPU call server-side, so throughput scales with concurrency: one RTX PRO 6000 sustains 200
+six-minute files per minute.
+
+Deploy from the [Model Library](https://app.baseten.co/deploy/baseten/nemotron-3-diarization-batch); see the [parent README](../README.md) for
+model details, accuracy and the other presets.
+
+## Endpoint
+
+```
+POST https://model-{MODEL_ID}.api.baseten.co/environments/production/predict
+Authorization: Api-Key {BASETEN_API_KEY}
+Content-Type: application/json
+```
+
+## Request
+
+```json
+{
+  "diarization_input": {
+    "audio": {"url": "https://github.com/ggerganov/whisper.cpp/raw/master/samples/jfk.wav"},
+    "latency": "offline"
+  }
+}
+```
+
+| Field | Type | Required | Description |
+|---|---|---|---|
+| `diarization_input` | object | yes | Wrapper object. A request without it is rejected with HTTP 400. |
+| `diarization_input.audio` | object | yes | Exactly one of `url` or `audio_b64`. |
+| `diarization_input.audio.url` | string | one of | Public or presigned URL. Fetched server-side with a 120 s timeout. |
+| `diarization_input.audio.audio_b64` | string | one of | Base64 of the audio file bytes (the whole file, not raw PCM). Validated as base64. |
+| `diarization_input.latency` | string | no | `offline` (default, 30.4 s buffer, best accuracy), `low` (1.04 s) or `ultralow` (0.32 s). |
+
+Audio may be any codec `ffmpeg` can decode from a stream (WAV, FLAC, MP3, OGG, WebM/Opus, …), any
+sample rate, mono or stereo; it is decoded and resampled to 16 kHz mono on the server. See Errors
+for the MP4/M4A caveat. Use FLAC or a
+compressed format for `audio_b64` to stay under Baseten's request-size limit on long files.
+
+## Response
+
+Real response for the 11-second `jfk.wav` sample above, `offline` profile (server time 936 ms, of
+which 527 ms was fetching the URL):
+
+```json
+{
+  "turns": [
+    {
+      "start": 0.29,
+      "end": 2.31,
+      "speaker": "speaker_0"
+    },
+    {
+      "start": 3.26,
+      "end": 4.56,
+      "speaker": "speaker_0"
+    },
+    {
+      "start": 5.37,
+      "end": 10.63,
+      "speaker": "speaker_0"
+    }
+  ],
+  "segments": [
+    "0.290 2.310 speaker_0",
+    "3.260 4.560 speaker_0",
+    "5.370 10.630 speaker_0"
+  ],
+  "num_speakers": 1,
+  "latency": "offline",
+  "batch_n": 1,
+  "timing": {
+    "fwd_ms": 7,
+    "h2d_ms": 0,
+    "mel_ms": 1,
+    "steps_ms": 5,
+    "graph_bs": 4,
+    "nonfinite_preds": 0,
+    "nonfinite_state": 0,
+    "pred_max": 1.0,
+    "pred_min": 0.0,
+    "core_dtype": "bfloat16",
+    "batch_ms": 7,
+    "batch_n": 1,
+    "batch_audio_s": 11.0,
+    "gap_ms": 1477,
+    "queue_ms": 401,
+    "decode_ms": 125,
+    "audio_s": 11.0,
+    "post_ms": 1,
+    "fetch_ms": 527,
+    "total_ms": 936
+  }
+}
+```
+
+| Field | Type | Description |
+|---|---|---|
+| `turns` | array | Speaker turns, ordered by start time. `start`/`end` in seconds from the beginning of the file; `speaker` is `speaker_0` … `speaker_7`, numbered in order of first appearance. Turns of different speakers may overlap. Zero-length turns are dropped. |
+| `turns[].start`, `turns[].end` | number | Seconds, 10 ms resolution. |
+| `turns[].speaker` | string | Session-local label, not an identity across files. |
+| `segments` | array of strings | The same turns in NeMo's `"start end speaker"` text form. |
+| `num_speakers` | integer | Number of distinct speakers found (1–8). |
+| `latency` | string | The profile that ran. |
+| `batch_n` | integer | How many concurrent requests shared this GPU call (1 when alone). |
+| `timing` | object | Server-side profile of this request, milliseconds. The useful ones: `fetch_ms` (URL download or base64 decode), `decode_ms` (ffmpeg), `queue_ms` (wait for the batch window), `fwd_ms` (model), `post_ms` (turn extraction), `total_ms`; `audio_s` is the decoded duration. The remaining keys (`h2d_ms`, `mel_ms`, `steps_ms`, `graph_bs`, `gap_ms`, `batch_audio_s`, `core_dtype`, `pred_*`, `nonfinite_*`) are engine diagnostics and may change between releases. |
+
+## Errors
+
+Client errors return HTTP 400 with a JSON body `{"error": "<message>"}`:
+
+| `error` | Cause |
+|---|---|
+| `request must be {'diarization_input': {'audio': {...}}}` | Missing wrapper or audio object. |
+| `audio requires 'url' or 'audio_b64'` | Neither given. |
+| `audio_b64 is not valid base64: …` | Bad encoding. |
+| `could not fetch audio.url: HTTP Error 406: Not Acceptable` (or 403/404/timeout) | The host refused the server's plain `GET` (some hosts gate on `User-Agent`), the object does not exist, or the download exceeded 120 s. Use a presigned object-store URL, or send the file as `audio_b64`. |
+| `audio could not be decoded: …` | `ffmpeg` rejected the bytes. Audio is decoded from a pipe, so **MP4/M4A/MOV containers whose index (`moov`) sits at the end of the file cannot be read** — remux them (`ffmpeg -movflags +faststart`) or send WAV/FLAC/MP3/OGG. |
+| `latency must be one of ['offline', 'low', 'ultralow']` | Unknown profile. |
+
+Server faults return HTTP 500 and are logged with a traceback.
+
+## Limits
+
+- Up to 8 speakers per file; more speakers are merged into the nearest existing label.
+- A request waits at most 30 minutes for a GPU slot and 10 minutes for post-processing before
+  failing; in practice a one-hour file returns in a few seconds.
+- The server batches at most 32 files per GPU call and waits up to 1.5 s to fill a batch, so the
+  first request in a quiet period pays up to ~100 ms of batch-window latency.
+
+## Clients
+
+- [`client.py`](client.py) — `python client.py --url https://github.com/ggerganov/whisper.cpp/raw/master/samples/jfk.wav` or `python client.py --file local.flac` (base64), with `--latency`.
+- [`curl.sh`](curl.sh) — `./curl.sh https://…/audio.wav [offline|low|ultralow]`.
+
+Both read `BASETEN_API_KEY` and `MODEL_ID` from the environment.

nemotron-3-diarization/batch/client.py

@@ -0,0 +1,60 @@
+"""Diarize a recording with the Nemotron 3 Diarization batch endpoint.
+
+    export BASETEN_API_KEY=...  MODEL_ID=...
+    python client.py --url https://example.com/meeting.wav
+    python client.py --file meeting.flac --latency low      # base64 upload
+
+Prints one line per speaker turn: start, end, speaker.
+"""
+
+import argparse
+import base64
+import os
+import sys
+
+import requests
+
+
+def diarize(audio: dict, latency: str = "offline") -> dict:
+    url = f"https://model-{os.environ['MODEL_ID']}.api.baseten.co/environments/production/predict"
+    resp = requests.post(
+        url,
+        headers={"Authorization": f"Api-Key {os.environ['BASETEN_API_KEY']}"},
+        json={"diarization_input": {"audio": audio, "latency": latency}},
+        timeout=600,
+    )
+    resp.raise_for_status()
+    return resp.json()
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    src = ap.add_mutually_exclusive_group(required=True)
+    src.add_argument("--url", help="public or presigned URL of the audio file")
+    src.add_argument(
+        "--file", help="local audio file, sent as base64 (any ffmpeg-decodable format)"
+    )
+    ap.add_argument(
+        "--latency", default="offline", choices=["offline", "low", "ultralow"]
+    )
+    args = ap.parse_args()
+
+    if args.url:
+        audio = {"url": args.url}
+    else:
+        with open(args.file, "rb") as f:
+            audio = {"audio_b64": base64.b64encode(f.read()).decode()}
+
+    result = diarize(audio, args.latency)
+    for turn in result["turns"]:
+        print(f"{turn['start']:8.2f}  {turn['end']:8.2f}  {turn['speaker']}")
+    t = result["timing"]
+    print(
+        f"\n{result['num_speakers']} speaker(s), {t['audio_s']:.1f} s of audio, "
+        f"profile={result['latency']}, server {t['total_ms']} ms",
+        file=sys.stderr,
+    )
+
+
+if __name__ == "__main__":
+    main()

nemotron-3-diarization/batch/curl.sh

@@ -0,0 +1,11 @@
+#!/usr/bin/env bash
+# Diarize a recording by URL with the Nemotron 3 Diarization batch endpoint.
+#   export BASETEN_API_KEY=... MODEL_ID=...
+#   ./curl.sh https://example.com/meeting.wav [offline|low|ultralow]
+set -euo pipefail
+URL=${1:?audio url}; LATENCY=${2:-offline}
+curl -sS -X POST "https://model-${MODEL_ID}.api.baseten.co/environments/production/predict" \
+  -H "Authorization: Api-Key ${BASETEN_API_KEY}" \
+  -H "Content-Type: application/json" \
+  -d "{\"diarization_input\": {\"audio\": {\"url\": \"${URL}\"}, \"latency\": \"${LATENCY}\"}}"
+echo