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.
This commit is contained in:
@@ -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())}
|
||||
Reference in New Issue
Block a user