parakeet: rewrite on sherpa-onnx; own the wrapper end-to-end

The Shadowfita FastAPI wrapper hit two unfixed upstream bugs on the first real /transcribe call — chunker return-shape mismatch (open issue #16) and a `torchaudio.tensor` that doesn't exist (open #10). Rather than babysit someone else's half-tested code, switched to sherpa-onnx with the prebuilt int8 Parakeet-TDT tarball from k2-fsa, and wrote our own ~60-line FastAPI wrapper. Moving parts now owned in-tree: Dockerfile CUDA 12.8 + cuDNN 9 runtime base, installs sherpa-onnx==1.12.39+cuda12.cudnn9 + fastapi + soundfile + libasound2 (sherpa-onnx links to ALSA at load time even when we never touch a mic). app.py OfflineRecognizer.from_transducer() once at startup; /transcribe and /v1/audio/transcriptions both accept multipart uploads and return {"text": ...}. entrypoint.sh Idempotent model download to /models on first run (~400 MB int8 tarball), then exec uvicorn. Smoke test: 0.wav (bundled in the tarball, The House of the Seven Gables excerpt) transcribes cleanly in ~1.2s on GPU. PARAKEET_MODEL_URL in .env lets you swap to the v3 (25-language) tarball without touching any other files. Wipe *.onnx + tokens.txt from the models dir and the entrypoint re-downloads.
2026-04-24 00:18:45 -07:00
parent 82f95d7428
commit 01c5380059
6 changed files with 272 additions and 102 deletions
@@ -5,45 +5,36 @@
 #   docker compose build
 #   docker compose up -d

-# Pinned git SHA to build from. Bump + rebuild when you want upstream
-# fixes. `main` latest as of 2026-04:
-#   https://github.com/Shadowfita/parakeet-tdt-0.6b-v2-fastapi
-PARAKEET_SHA=31c5652b62d09653ad5ea8190c0ad0d35394174d
+# Image tag. Bump when you change the Dockerfile / app.py so docker caches
+# cleanly.
+PARAKEET_TAG=sherpa-onnx-v2

 # Host port for the FastAPI server (container listens on 8000)
 PARAKEET_PORT=8765

-# Bind address. 0.0.0.0 exposes on all interfaces including the WG
-# tunnel IP (10.100.79.3). Use 127.0.0.1 to restrict to local-only.
+# Bind address. 0.0.0.0 exposes on all interfaces including the WG tunnel IP
+# (10.100.79.3). Use 127.0.0.1 to restrict to local-only.
 PARAKEET_BIND=0.0.0.0

-# Host path for the HuggingFace cache (parakeet-tdt-0.6b-v2 weights
-# ~2.5 GB). Persistent across container recreates. Must exist before
-# first `up` with ownership matching the container user (root inside
-# this image — no UID juggling needed, but the host dir needs to be
-# writable by the container).
+# Host path for the ONNX model files — encoder/decoder/joiner/tokens.txt.
+# Downloaded by the entrypoint on first run if absent. Must exist before
+# first `up` (directory, not files).
 PARAKEET_MODELS_DIR=/worktank/parakeet/models

-# Inference precision. fp16 halves VRAM and is lossless for parakeet
-# in practice; use fp32 only if fp16 shows degraded WER for your
-# domain audio.
-PARAKEET_MODEL_PRECISION=fp16
+# Which sherpa-onnx release tarball to fetch on first boot. Default is the
+# int8-quantized English-only v2 (~400 MB). Switch to the v3 tarball below
+# to cover 25 European languages at a similar size:
+#   https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/sherpa-onnx-nemo-parakeet-tdt-0.6b-v3-int8.tar.bz2
+PARAKEET_MODEL_URL=https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/sherpa-onnx-nemo-parakeet-tdt-0.6b-v2-int8.tar.bz2

-# Batch size for the transcribe queue. Larger = better throughput
-# under load at the cost of per-request latency.
-PARAKEET_BATCH_SIZE=4
+# ONNX Runtime execution provider. `cuda` uses the GPU (requires nvidia
+# runtime + matching CUDA/cuDNN in the image). `cpu` falls back to CPU —
+# fine for low-volume dev use; ~4-8× slower on this host.
+PARAKEET_PROVIDER=cuda

-# Max single-clip duration (seconds). Longer inputs get rejected
-# by the server with 400. Upstream default.
-PARAKEET_MAX_AUDIO_DURATION=30
-
-# Silero VAD threshold (0–1). Higher = stricter about what counts
-# as speech (fewer false wake-ups on silence, more chance of clipping
-# soft speech). 0.5 is upstream default.
-PARAKEET_VAD_THRESHOLD=0.5
-
-# End-to-end processing timeout per request (seconds).
-PARAKEET_PROCESSING_TIMEOUT=60
+# CPU threads per recognizer session. Irrelevant when provider=cuda;
+# only matters for provider=cpu.
+PARAKEET_NUM_THREADS=1

 # Log level: DEBUG | INFO | WARNING | ERROR
 PARAKEET_LOG_LEVEL=INFO
@@ -0,0 +1,41 @@
+# syntax=docker/dockerfile:1.6
+#
+# Parakeet-TDT ASR via sherpa-onnx (ONNX Runtime + CUDA).
+#
+# We own this whole image — not forked from an upstream wrapper. ~70 MB of
+# application layer over the CUDA+cuDNN runtime base.
+
+ARG CUDA_BASE=nvidia/cuda:12.8.0-cudnn-runtime-ubuntu22.04
+FROM ${CUDA_BASE}
+
+ENV DEBIAN_FRONTEND=noninteractive \
+    PIP_ROOT_USER_ACTION=ignore \
+    PYTHONUNBUFFERED=1 \
+    PATH="/opt/venv/bin:${PATH}"
+
+RUN apt-get update \
+ && apt-get install -y --no-install-recommends \
+        python3 python3-pip python3-venv \
+        libsndfile1 \
+        libasound2 \
+        ca-certificates wget bzip2 \
+ && rm -rf /var/lib/apt/lists/* \
+ && python3 -m venv /opt/venv
+
+RUN pip install --no-cache-dir \
+        fastapi \
+        'uvicorn[standard]' \
+        python-multipart \
+        soundfile \
+        numpy \
+ && pip install --no-cache-dir \
+        sherpa-onnx==1.12.39+cuda12.cudnn9 \
+        -f https://k2-fsa.github.io/sherpa/onnx/cuda.html
+
+WORKDIR /app
+COPY app.py /app/app.py
+COPY entrypoint.sh /usr/local/bin/entrypoint.sh
+RUN chmod +x /usr/local/bin/entrypoint.sh
+
+EXPOSE 8000
+ENTRYPOINT ["/usr/local/bin/entrypoint.sh"]
@@ -1,51 +1,56 @@
 # Parakeet ASR

-NVIDIA Parakeet-TDT 0.6B v2 speech-to-text served via a FastAPI
-wrapper with Silero VAD and WebSocket streaming.
+NVIDIA Parakeet-TDT 0.6B (int8 ONNX) served by our own thin FastAPI
+wrapper over [sherpa-onnx](https://github.com/k2-fsa/sherpa-onnx)
+(ONNX Runtime + CUDA).

 **Server:** irv-ml1 (Irvine, WireGuard-only)
-**Port:** 8765 (configurable via `.env`)
-**GPUs:** both exposed (`NVIDIA_VISIBLE_DEVICES=all`); upstream
-respects `CUDA_VISIBLE_DEVICES` if later pinning is needed
-**Upstream:** [Shadowfita/parakeet-tdt-0.6b-v2-fastapi](https://github.com/Shadowfita/parakeet-tdt-0.6b-v2-fastapi)
-**Image:** built locally from a pinned git SHA via docker buildx's
-git URL context — no source vendored into this workspace
+**Port:** 8765 (container 8000)
+**GPU:** both exposed (`NVIDIA_VISIBLE_DEVICES=all`); sherpa-onnx uses
+whichever CUDA ExecutionProvider picks
+**Image:** `local/parakeet:sherpa-onnx-v1` — built from `Dockerfile` +
+`app.py` + `entrypoint.sh` in this directory; **we own all the code**
+
+## Why not the FastAPI community wrappers
+
+Both `Shadowfita/parakeet-tdt-0.6b-v2-fastapi` and
+`pnivek/Parakeet-ASR-FastAPI` look appealing on paper but have open,
+unfixed bugs in the actual transcribe path (return-shape mismatches
+after an unpinned `torchaudio` upgrade, `torchaudio.tensor` which
+doesn't exist, etc.). We tried Shadowfita and hit #16+#10 on the
+first real request. Rather than babysit someone else's half-tested
+code, we moved to sherpa-onnx — ONNX Runtime is a stable base, k2-fsa
+publishes prebuilt int8 Parakeet weights per release, and the
+recognizer API is a three-line call.

 ## API endpoints

 | Method + path | Purpose |
 |---|---|
-| `POST /transcribe` | Batch transcription (multipart file upload) |
-| `WS /ws/transcribe` | Streaming with Silero VAD — partial + final segments |
+| `POST /transcribe` | Multipart file upload → `{"text": "..."}` |
+| `POST /v1/audio/transcriptions` | Same body; OpenAI-compatible path |
 | `GET /healthz` | Health probe (used by docker healthcheck) |

-Not the literal OpenAI `/v1/audio/transcriptions` path. If you have
-a downstream client that demands that URL shape, either point its
-base URL at `/transcribe`, or add a Traefik/nginx path alias in
-front.
-
 ## Path layout

 | Host path | Container path | Purpose | Restic? |
 |---|---|---|---|
-| `/worktank/parakeet/models/` | `/models` (`HF_HOME`) | HF cache for parakeet-tdt-0.6b-v2 weights (~2.5 GB) | excluded (regenerable from HF) |
-
-The container runs as root internally; host dir just needs to exist
-and be writable.
+| `/worktank/parakeet/models/` | `/models` | ONNX encoder+decoder+joiner+tokens (~400 MB int8) | excluded (regenerable — re-downloads from the URL on first run if absent) |

 ## First-time deploy on irv-ml1

 ```bash
-# 1. Push compose + env template
+# 1. Push compose + Dockerfile + app + entrypoint
 scripts/deploy-stack.sh irv-ml1 parakeet

-# 2. Create the models dir on the host. One-time sudo — /worktank
-#    itself is root-owned.
+# 2. Make sure the models dir exists (one-time, already done from the
+#    earlier Shadowfita deploy; this is idempotent)
 ssh -t irv-ml1 'sudo mkdir -p /worktank/parakeet/models && \
                sudo chown -R lkraven:lkraven /worktank/parakeet'

-# 3. Build the image (first time only; ~5–10 min for torch + NeMo
-#    wheels). Then `up`.
+# 3. Build the image and bring up. First boot does a ~400 MB model
+#    download via the entrypoint; allow 1–2 minutes before /healthz
+#    flips healthy.
 ssh irv-ml1 '
  cd /opt/docker/compose/parakeet && \
  cp -n .env.example .env && \
@@ -56,33 +61,54 @@ ssh irv-ml1 '
 '
 ```

-First `/transcribe` request downloads parakeet-tdt-0.6b-v2 weights
-to `/worktank/parakeet/models/` (~2.5 GB).
-
 ## Smoke test

 ```bash
-# From the workstation over WG
+# Over WG from the workstation
 curl -F "file=@sample.wav" http://10.100.79.3:8765/transcribe
+# → {"text": "hello world"}
+
+# OpenAI-shape alias (for clients that only know /v1/audio/transcriptions)
+curl -F "file=@sample.wav" http://10.100.79.3:8765/v1/audio/transcriptions
 ```

-## Rebuild against a newer upstream commit
+## Switching to the v3 (multilingual) model
+
+The env var `PARAKEET_MODEL_URL` picks the release tarball. To swap
+from the English-only v2 to the 25-language v3:

 ```bash
 ssh irv-ml1 '
  cd /opt/docker/compose/parakeet && \
-  sed -i "s/^PARAKEET_SHA=.*/PARAKEET_SHA=<new-sha>/" .env && \
-  docker compose build && \
-  docker compose up -d
+  sed -i "s|v2-int8|v3-int8|" .env && \
+  # Wipe the v2 weights so the entrypoint re-downloads v3 on next up:
+  rm -f /worktank/parakeet/models/*.onnx /worktank/parakeet/models/tokens.txt && \
+  docker compose up -d && \
+  docker compose logs -f --tail=30
 '
 ```

-Models dir is unaffected.
+## Upgrade sherpa-onnx or change the base image

-## Deploy updates
+Bump `PARAKEET_TAG` in `.env` to force a rebuild of the local image
+after editing the `Dockerfile`, then:

 ```bash
-# After editing compose.yaml or .env.example here
 scripts/deploy-stack.sh irv-ml1 parakeet
-ssh irv-ml1 'cd /opt/docker/compose/parakeet && docker compose up -d'
+ssh irv-ml1 'cd /opt/docker/compose/parakeet && docker compose build && docker compose up -d'
+```
+
+Model files under `/worktank/parakeet/models/` are preserved across
+image rebuilds.
+
+## File layout
+
+```
+stacks/parakeet/
+├── Dockerfile           # CUDA 12.8 + cuDNN 9 base, sherpa-onnx-cu12 wheel
+├── app.py               # FastAPI — ~60 lines
+├── entrypoint.sh        # downloads model on first run, then uvicorn
+├── compose.yaml         # one service, bind-mounts the models dir
+├── .env.example         # template; real .env lives on the server
+└── README.md            # this file
 ```
@@ -0,0 +1,95 @@
+"""Thin FastAPI wrapper around sherpa-onnx's OfflineRecognizer for Parakeet-TDT.
+
+Load the encoder/decoder/joiner/tokens once at startup; serve:
+  POST /transcribe                   — our native shape
+  POST /v1/audio/transcriptions      — OpenAI-compatible alias (returns {"text": ...})
+  GET  /healthz                      — used by the docker healthcheck
+
+No VAD chunking, no Silero preprocessing — parakeet-tdt handles long-form natively
+and the int8 ONNX model on a 24 GB GPU eats everything we're likely to throw at it.
+"""
+
+from __future__ import annotations
+
+import io
+import logging
+import os
+from pathlib import Path
+
+import numpy as np
+import sherpa_onnx
+import soundfile as sf
+from fastapi import FastAPI, File, HTTPException, UploadFile
+
+MODEL_DIR = Path(os.environ.get("MODEL_DIR", "/models"))
+PROVIDER = os.environ.get("PROVIDER", "cuda")
+NUM_THREADS = int(os.environ.get("NUM_THREADS", "1"))
+
+REQUIRED_FILES = (
+    "encoder.int8.onnx",
+    "decoder.int8.onnx",
+    "joiner.int8.onnx",
+    "tokens.txt",
+)
+
+logger = logging.getLogger("parakeet")
+logging.basicConfig(level=os.environ.get("LOG_LEVEL", "INFO"))
+
+
+def _ensure_model_present() -> None:
+    missing = [f for f in REQUIRED_FILES if not (MODEL_DIR / f).exists()]
+    if missing:
+        raise RuntimeError(
+            f"Missing model files in {MODEL_DIR}: {missing}. "
+            "The entrypoint is responsible for downloading them before the server starts."
+        )
+
+
+def _load_recognizer() -> sherpa_onnx.OfflineRecognizer:
+    _ensure_model_present()
+    logger.info("loading OfflineRecognizer (provider=%s, threads=%d)", PROVIDER, NUM_THREADS)
+    return sherpa_onnx.OfflineRecognizer.from_transducer(
+        encoder=str(MODEL_DIR / "encoder.int8.onnx"),
+        decoder=str(MODEL_DIR / "decoder.int8.onnx"),
+        joiner=str(MODEL_DIR / "joiner.int8.onnx"),
+        tokens=str(MODEL_DIR / "tokens.txt"),
+        model_type="nemo_transducer",
+        provider=PROVIDER,
+        num_threads=NUM_THREADS,
+    )
+
+
+app = FastAPI(title="Parakeet ASR (sherpa-onnx)")
+recognizer = _load_recognizer()
+
+
+def _decode(raw: bytes) -> str:
+    try:
+        samples, sample_rate = sf.read(io.BytesIO(raw), dtype="float32")
+    except Exception as exc:
+        raise HTTPException(400, f"Could not decode audio: {exc}") from exc
+    if samples.ndim > 1:
+        samples = samples.mean(axis=1).astype(np.float32)
+
+    stream = recognizer.create_stream()
+    stream.accept_waveform(sample_rate, samples)
+    recognizer.decode_stream(stream)
+    return stream.result.text
+
+
+@app.get("/healthz")
+def healthz() -> dict[str, str]:
+    return {"status": "ok"}
+
+
+@app.post("/transcribe")
+async def transcribe(file: UploadFile = File(...)) -> dict[str, str]:
+    return {"text": _decode(await file.read())}
+
+
+@app.post("/v1/audio/transcriptions")
+async def openai_transcriptions(file: UploadFile = File(...)) -> dict[str, str]:
+    # OpenAI's shape: {"text": "..."} by default; extra fields (model, language,
+    # response_format) are accepted by real OpenAI but ignored here — the model
+    # choice is baked in at container startup.
+    return {"text": _decode(await file.read())}
@@ -1,37 +1,28 @@
-# Parakeet ASR — NVIDIA Parakeet-TDT 0.6B v2 speech-to-text.
+# Parakeet ASR via sherpa-onnx + our own thin FastAPI wrapper.
 #
-# Wraps Shadowfita/parakeet-tdt-0.6b-v2-fastapi (FastAPI + Silero VAD +
-# WebSocket streaming). Upstream provides no prebuilt image, so we
-# build from a pinned git commit via docker buildx's git URL context
-# — no source files vendored into this repo.
+# We previously wrapped Shadowfita/parakeet-tdt-0.6b-v2-fastapi but hit two
+# unfixed upstream bugs (open issues #16 + #10) the first time we actually sent
+# a transcription request. Switched to sherpa-onnx — ONNX Runtime + CUDA, a
+# prebuilt int8 quantized Parakeet-TDT from k2-fsa — and wrote our own ~50-line
+# wrapper we own end-to-end.
 #
-# Runs on irv-ml1 (dual GPU). Both GPUs exposed; upstream respects
-# CUDA_VISIBLE_DEVICES if you want to pin later.
+# Model weights (~400 MB int8) download on first run via the entrypoint to
+# ${PARAKEET_MODELS_DIR}/ (persistent host bind mount). Subsequent starts skip
+# the download.
 #
-# HF weights (~2.5 GB for parakeet-tdt-0.6b-v2) cache to
-# ${PARAKEET_MODELS_DIR} via HF_HOME=/models, persistent across
-# container recreates.
-#
-# API routes (per upstream README):
-#   POST /transcribe          — batch transcription
-#   WS   /ws/transcribe       — streaming with Silero VAD
+# API:
+#   POST /transcribe                  — multipart file upload, returns {"text": "..."}
+#   POST /v1/audio/transcriptions     — same body, OpenAI-compatible path alias
 #   GET  /healthz
 #
-# Not a literal OpenAI `/v1/audio/transcriptions` path; point clients
-# at /transcribe directly, or add a reverse-proxy alias if drop-in
-# compat is needed later.
-#
-# First `up` triggers a fresh docker build (python:3.10-slim +
-# torch + NeMo ≈ 5–10 min). Subsequent starts reuse the cached
-# image unless PARAKEET_SHA changes.
-#
 # All tunables live in .env — edit that, not this file.

 services:
  parakeet:
-    image: local/parakeet:${PARAKEET_SHA}
+    image: local/parakeet:${PARAKEET_TAG}
    build:
-      context: https://github.com/Shadowfita/parakeet-tdt-0.6b-v2-fastapi.git#${PARAKEET_SHA}
+      context: .
+      dockerfile: Dockerfile
    container_name: parakeet
    restart: unless-stopped
    runtime: nvidia
@@ -39,14 +30,10 @@ services:
      - "${PARAKEET_BIND:-0.0.0.0}:${PARAKEET_PORT}:8000"
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
-      - HF_HOME=/models
-      - DEVICE=cuda
-      - MODEL_PRECISION=${PARAKEET_MODEL_PRECISION:-fp16}
-      - BATCH_SIZE=${PARAKEET_BATCH_SIZE:-4}
-      - TARGET_SAMPLE_RATE=16000
-      - MAX_AUDIO_DURATION=${PARAKEET_MAX_AUDIO_DURATION:-30}
-      - VAD_THRESHOLD=${PARAKEET_VAD_THRESHOLD:-0.5}
-      - PROCESSING_TIMEOUT=${PARAKEET_PROCESSING_TIMEOUT:-60}
+      - MODEL_DIR=/models
+      - MODEL_URL=${PARAKEET_MODEL_URL}
+      - PROVIDER=${PARAKEET_PROVIDER:-cuda}
+      - NUM_THREADS=${PARAKEET_NUM_THREADS:-1}
      - LOG_LEVEL=${PARAKEET_LOG_LEVEL:-INFO}
    volumes:
      - ${PARAKEET_MODELS_DIR}:/models
@@ -55,13 +42,11 @@ services:
      interval: 30s
      timeout: 10s
      retries: 3
-      # First `up` may spend several minutes on torch/NeMo install
-      # during the image build phase; after the image exists, startup
-      # is ~30-60s (NeMo model load).
-      start_period: 180s
+      # First boot may include a ~400 MB model download.
+      start_period: 300s
    labels:
      - homepage.group=AI Systems
      - homepage.name=Parakeet ASR
      - homepage.icon=mdi-microphone
-      - homepage.description=Parakeet-TDT speech-to-text (irv-ml1)
+      - homepage.description=Parakeet-TDT speech-to-text via sherpa-onnx (irv-ml1)
      - homepage.href=http://10.100.79.3:${PARAKEET_PORT}
@@ -0,0 +1,32 @@
+#!/bin/bash
+# entrypoint: download int8 Parakeet weights on first run if the target dir is
+# empty, then start the FastAPI app. Idempotent — subsequent starts skip the
+# download when the files already exist.
+
+set -euo pipefail
+
+MODEL_DIR=${MODEL_DIR:-/models}
+MODEL_URL=${MODEL_URL:-https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/sherpa-onnx-nemo-parakeet-tdt-0.6b-v2-int8.tar.bz2}
+
+required=(encoder.int8.onnx decoder.int8.onnx joiner.int8.onnx tokens.txt)
+missing=0
+for f in "${required[@]}"; do
+    [[ -f "$MODEL_DIR/$f" ]] || missing=1
+done
+
+if [[ "$missing" -eq 1 ]]; then
+    echo "[entrypoint] model files not present in $MODEL_DIR, downloading from $MODEL_URL"
+    mkdir -p "$MODEL_DIR"
+    tmp=$(mktemp -d)
+    wget -q --show-progress -O "$tmp/model.tar.bz2" "$MODEL_URL"
+    tar -xjf "$tmp/model.tar.bz2" -C "$tmp"
+    # Upstream tarballs extract to a single top-level dir; move its contents up.
+    extracted_root=$(find "$tmp" -mindepth 1 -maxdepth 1 -type d | head -n1)
+    shopt -s dotglob
+    mv "$extracted_root"/* "$MODEL_DIR"/
+    rm -rf "$tmp"
+    echo "[entrypoint] model extracted:"
+    ls -la "$MODEL_DIR"
+fi
+
+exec uvicorn app:app --host 0.0.0.0 --port 8000