nvidia parakeet-tdt-0.6b-v3 openai realtime asr

nvasr.py

import asyncio
import base64
import json
import logging
import time
import uuid
import numpy as np
import torch
import uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import JSONResponse
from nemo.collections.asr.models import ASRModel

# ---------------------------------------------------------
# 0. Logging Configuration
# ---------------------------------------------------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-7s | %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger("parakeet-realtime-asr")

PORT = 5551
SAMPLE_RATE = 16000

app = FastAPI(title="Parakeet Realtime ASR - whisrs / OpenAI Realtime Compatible")

# ---------------------------------------------------------
# 1. Model Initialization
# ---------------------------------------------------------
if torch.cuda.is_available():
    device_str = "cuda:0"
    logger.info(f"CUDA available: {torch.cuda.get_device_name(0)}")
else:
    device_str = "cpu"
    logger.warning("!!! CUDA NOT AVAILABLE - running on CPU !!!")

logger.info(f"Loading nvidia/parakeet-tdt-0.6b-v2 on {device_str}...")
asr_model = ASRModel.from_pretrained("nvidia/parakeet-tdt-0.6b-v2")
asr_model = asr_model.to(device_str)
asr_model.eval()

if device_str.startswith("cuda"):
    try:
        asr_model.half()
        logger.info("Model converted to fp16.")
    except Exception as e:
        logger.warning(f"fp16 conversion failed, using fp32: {e}")

logger.info("Parakeet model loaded and ready.")

# ---------------------------------------------------------
# 2. Health / Service Discovery Endpoint
# ---------------------------------------------------------
@app.get("/api/v1/health")
@app.get("/health")
async def health_check():
    response = {
        "status": "healthy",
        "model": "nvidia/parakeet-tdt-0.6b-v2",
        "websocket_port": PORT,
        "device": device_str,
    }
    logger.info(f"<<< [SENT] Health response: {response}")
    return JSONResponse(content=response)

# ---------------------------------------------------------
# 3. Helpers
# ---------------------------------------------------------
def generate_id(prefix: str = "evt") -> str:
    return f"{prefix}_{uuid.uuid4().hex[:12]}"

def detect_and_decode(raw: bytes) -> tuple[np.ndarray, str]:
    """Try common PCM layouts (each truncated to whole elements) and pick highest RMS."""
    n = len(raw)
    logger.info(f"    [DETECT] Buffer: {n} bytes (n%2={n%2}, n%4={n%4})")

    if n == 0:
        return np.array([], dtype=np.float32), "empty"

    if raw.count(0) == n:
        logger.warning("    [DETECT] ENTIRE buffer is zero bytes -> client is sending silence!")

    candidates: dict[str, np.ndarray] = {}

    if n >= 2:
        raw16 = raw[: (n // 2) * 2]
        candidates["s16le"] = np.frombuffer(raw16, "<i2").astype(np.float32) / 32768.0
    if n >= 4:
        raw32 = raw[: (n // 4) * 4]
        candidates["f32le"] = np.frombuffer(raw32, "<f4").astype(np.float32)
        candidates["s32le"] = np.frombuffer(raw32, "<i4").astype(np.float32) / 2147483648.0
        candidates["s32le_hi16"] = (
            np.frombuffer(raw32, "<i4").astype(np.int64) >> 16
        ).astype(np.float32) / 32768.0

    best_name, best_arr = "s16le", np.array([], dtype=np.float32)
    best_rms = 0.0
    for name, arr in candidates.items():
        rms = float(np.sqrt(np.mean(arr.astype(np.float64) ** 2))) if len(arr) else 0.0
        nz = float(np.count_nonzero(arr)) / max(len(arr), 1)
        logger.info(f"    [DETECT] {name:9s} rms={rms:.6f} nonzero={nz*100:.1f}% samples={len(arr)}")
        if rms > best_rms:
            best_rms, best_name, best_arr = rms, name, arr

    logger.info(f"    [DETECT] Selected format: {best_name} (rms={best_rms:.6f})")
    return best_arr.astype(np.float32), best_name

def run_inference(audio_array: np.ndarray) -> str:
    """Synchronous inference executed inside an async worker thread."""
    if len(audio_array) == 0:
        logger.info("    [ANALYZE] Empty audio array, skipping inference")
        return ""

    duration_sec = len(audio_array) / float(SAMPLE_RATE)
    rms = float(np.sqrt(np.mean(audio_array.astype(np.float64) ** 2)))
    logger.info(
        f"    [ANALYZE] Inference on {len(audio_array)} samples "
        f"({duration_sec:.2f}s @ {SAMPLE_RATE}Hz) | device="
        f"{next(asr_model.parameters()).device} | rms={rms:.4f}"
    )

    t0 = time.perf_counter()
    with torch.no_grad():
        transcriptions = asr_model.transcribe([audio_array], batch_size=1)
    t1 = time.perf_counter()

    # NeMo may return Hypothesis objects or plain strings depending on version
    raw_out = transcriptions[0] if transcriptions else ""
    if hasattr(raw_out, "text"):
        transcript = raw_out.text.strip()
    else:
        transcript = str(raw_out).strip()

    rtfx = duration_sec / max(t1 - t0, 1e-6)
    logger.info(
        f"    [ANALYZE] Done in {(t1 - t0)*1000:.1f}ms (RTFx: {rtfx:.1f}x) -> {transcript!r}"
    )
    return transcript

# ---------------------------------------------------------
# 4. WebSocket Realtime Endpoint
# ---------------------------------------------------------
@app.websocket("/v1/realtime")
async def realtime_endpoint(websocket: WebSocket):
    await websocket.accept()
    session_id = generate_id("sess")
    client_host = f"{websocket.client.host}:{websocket.client.port}"
    logger.info(f"=== [CONNECTED] {client_host} (session: {session_id}) ===")

    audio_buffer = bytearray()

    # Initial session.created event
    await websocket.send_json({
        "event_id": generate_id("evt"),
        "type": "session.created",
        "session": {
            "id": session_id,
            "model": "nvidia/parakeet-tdt-0.6b-v2",
            "input_audio_format": "pcm16",
            "input_audio_transcription": {
                "model": "nvidia/parakeet-tdt-0.6b-v2"
            },
        },
    })
    logger.info("<<< [SENT] session.created")

    try:
        while True:
            try:
                message = await websocket.receive()
            except RuntimeError:
                logger.info("=== [DISCONNECTED] (RuntimeError) ===")
                break
            except WebSocketDisconnect:
                logger.info(f"=== [DISCONNECTED] {client_host} ===")
                break

            # -------------------------------------------------
            # Case A: Binary frames (raw PCM)
            # -------------------------------------------------
            if "bytes" in message and message["bytes"]:
                raw_bytes = message["bytes"]
                audio_buffer.extend(raw_bytes)
                logger.debug(
                    f">>> [RECEIVED BINARY] {len(raw_bytes)} bytes | buffer: {len(audio_buffer)} bytes"
                )

            # -------------------------------------------------
            # Case B: Text frames (JSON events)
            # -------------------------------------------------
            elif "text" in message and message["text"]:
                data = json.loads(message["text"])
                event_type = data.get("type", "unknown")
                logger.info(f">>> [RECEIVED EVENT] type={event_type!r}")

                # 1. Base64 audio chunk append  <-- WAS MISSING!
                if event_type == "input_audio_buffer.append":
                    base64_audio = data.get("audio", "")
                    if base64_audio:
                        raw_bytes = base64.b64decode(base64_audio)
                        audio_buffer.extend(raw_bytes)
                        logger.debug(f"    [BUFFER] +{len(raw_bytes)} bytes -> total {len(audio_buffer)}")
                    else:
                        logger.warning(">>> append event with empty 'audio' field")

                # 2. Commit -> run transcription
                elif event_type == "input_audio_buffer.commit":
                    item_id = generate_id("item")

                    await websocket.send_json({
                        "event_id": generate_id("evt"),
                        "type": "input_audio_buffer.committed",
                        "item_id": item_id,
                    })
                    logger.info(f"<<< [SENT] input_audio_buffer.committed (item {item_id})")

                    if len(audio_buffer) > 0:
                        raw_audio_bytes = bytes(audio_buffer)
                        audio_buffer.clear()

                        logger.info(
                            f"    [ANALYZE] Buffer: {len(raw_audio_bytes)} bytes | "
                            f"first 32 bytes hex: {raw_audio_bytes[:32].hex(' ')}"
                        )

                        audio_float32, fmt = detect_and_decode(raw_audio_bytes)

                        if len(audio_float32) == 0:
                            logger.warning("    [ANALYZE] Decoded array empty; sending empty transcript")
                            transcript = ""
                        else:
                            transcript = await asyncio.to_thread(run_inference, audio_float32)

                        # Send BOTH event names for whisrs + OpenAI compatibility
                        for evt_type in (
                            "input_audio_transcription.completed",
                            "conversation.item.input_audio_transcription.completed",
                        ):
                            await websocket.send_json({
                                "event_id": generate_id("evt"),
                                "type": evt_type,
                                "item_id": item_id,
                                "transcript": transcript,
                            })
                            logger.info(f"<<< [SENT] {evt_type} | transcript: {transcript!r}")
                    else:
                        logger.warning("    [BUFFER] Commit with empty buffer")
                        for evt_type in (
                            "input_audio_transcription.completed",
                            "conversation.item.input_audio_transcription.completed",
                        ):
                            await websocket.send_json({
                                "event_id": generate_id("evt"),
                                "type": evt_type,
                                "item_id": item_id,
                                "transcript": "",
                            })

                # 3. Clear buffer
                elif event_type == "input_audio_buffer.clear":
                    logger.info(f"    [BUFFER] Cleared {len(audio_buffer)} bytes")
                    audio_buffer.clear()
                    await websocket.send_json({
                        "event_id": generate_id("evt"),
                        "type": "input_audio_buffer.cleared",
                    })
                    logger.info("<<< [SENT] input_audio_buffer.cleared")

                # 4. Session update  <-- WAS MISSING! (likely caused the 15s timeout disconnect)
                elif event_type == "session.update":
                    session_cfg = data.get("session", {})
                    logger.info(f"    [SESSION] Update: {json.dumps(session_cfg)[:300]}")
                    await websocket.send_json({
                        "event_id": generate_id("evt"),
                        "type": "session.updated",
                        "session": session_cfg,
                    })
                    logger.info("<<< [SENT] session.updated")

                # 5. Ping / keepalive
                elif event_type == "ping":
                    await websocket.send_json({"type": "pong"})
                    logger.info("<<< [SENT] pong")

                else:
                    logger.warning(f"    [IGNORED] Unhandled event type: {event_type!r}")

    except WebSocketDisconnect:
        logger.info(f"=== [DISCONNECTED] {client_host} (session: {session_id}) ===")
    except Exception as e:
        logger.exception(f"Unexpected error in WebSocket loop: {e}")
        try:
            await websocket.send_json({
                "event_id": generate_id("evt"),
                "type": "error",
                "error": {"message": str(e), "type": "server_error"},
            })
            await websocket.close()
        except Exception:
            pass

    logger.info(f"=== [SESSION ENDED] {client_host} (session: {session_id}) ===")

# ---------------------------------------------------------
# 5. Entrypoint
# ---------------------------------------------------------
if __name__ == "__main__":
    logger.info(f"Starting server on 0.0.0.0:{PORT} (Device: {device_str})")
    uvicorn.run(app, host="0.0.0.0", port=PORT, log_level="warning")