omnius 1.0.0

package/dist/scripts/live-nemotron.py

@@ -0,0 +1,478 @@
+#!/usr/bin/env python3
+"""
+live-nemotron.py — Self-contained streaming ASR worker using NVIDIA's
+nvidia/nemotron-speech-streaming-en-0.6b model.
+
+Parallel to live-whisper.py. Same stdin/stdout protocol so the same
+pipelines (nexus voice subsystem, asr_listen tool, eval harness) can
+swap backends by pointing at a different script.
+
+Protocol:
+  stdin   — raw PCM16 (16kHz, mono, 16-bit signed little-endian)
+  stdout  — JSON lines:
+    {"type":"status","message":"Creating venv..."}
+    {"type":"status","message":"Installing dependencies..."}
+    {"type":"status","message":"Loading model..."}
+    {"type":"ready"}
+    {"type":"transcript","text":"hello world","isFinal":false}
+    {"type":"transcript","text":"hello world how are you","isFinal":true}
+    {"type":"error","message":"..."}
+
+Usage:
+  # Live stream from mic:
+  arecord -f S16_LE -r 16000 -c 1 -t raw -q - | python3 live-nemotron.py
+  # Single file transcription (write path + read transcript JSON):
+  python3 live-nemotron.py --file recording.wav
+
+Backend selection:
+  1. NeMo toolkit (nvidia NeMo) — native streaming support for Parakeet-
+     style models. Preferred when available.
+  2. transformers + torchaudio — fallback via HuggingFace's generic
+     ASR pipeline. Works for file-based transcription even when NeMo
+     install fails (common on macOS / no-CUDA setups). Does NOT do
+     streaming — buffers the full window each chunk.
+"""
+
+import sys
+import os
+import json
+import subprocess
+import struct
+import time
+import threading
+import argparse
+from pathlib import Path
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+VENV = SCRIPT_DIR / ".nemotron-venv"
+PY = VENV / "bin" / "python"
+PIP = VENV / "bin" / "pip"
+
+SAMPLE_RATE = 16000
+CHANNELS = 1
+SAMPLE_WIDTH = 2  # 16-bit
+CHUNK_SECONDS = 2.0    # Nemotron is a streaming model — shorter chunks than whisper
+WINDOW_SECONDS = 8.0
+
+# HuggingFace model identifier
+MODEL_ID = "nvidia/nemotron-speech-streaming-en-0.6b"
+
+# ---------------------------------------------------------------------------
+# Output helpers (JSON lines to stdout)
+# ---------------------------------------------------------------------------
+
+def emit(event: dict):
+    sys.stdout.write(json.dumps(event) + "\n")
+    sys.stdout.flush()
+
+
+def emit_status(msg: str):
+    emit({"type": "status", "message": msg})
+
+
+def emit_error(msg: str):
+    emit({"type": "error", "message": msg})
+
+
+def emit_transcript(text: str, is_final: bool = False, backend: str = "nemotron"):
+    emit({"type": "transcript", "text": text, "isFinal": is_final, "backend": backend})
+
+# ---------------------------------------------------------------------------
+# Venv bootstrap (same pattern as live-whisper.py)
+# ---------------------------------------------------------------------------
+
+def _in_venv() -> bool:
+    return sys.prefix != sys.base_prefix and str(SCRIPT_DIR) in sys.prefix
+
+
+def _ensure_venv():
+    if VENV.exists():
+        return
+    emit_status("Creating Python venv for Nemotron ASR...")
+    import venv
+    venv.EnvBuilder(with_pip=True).create(str(VENV))
+    subprocess.check_call(
+        [str(PY), "-m", "pip", "install", "--upgrade", "pip", "wheel"],
+        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
+    )
+
+
+def _ensure_deps():
+    """Install torch + either nemo_toolkit[asr] or transformers as fallback."""
+    need = []
+    try:
+        import numpy  # noqa: F401
+    except ImportError:
+        need.append("numpy")
+    try:
+        import torch  # noqa: F401
+    except ImportError:
+        need.append("torch")
+    try:
+        import soundfile  # noqa: F401
+    except ImportError:
+        need.append("soundfile")
+    try:
+        import transformers  # noqa: F401
+    except ImportError:
+        need.append("transformers")
+
+    if need:
+        emit_status(f"Installing core deps: {', '.join(need)}...")
+        try:
+            subprocess.check_call(
+                [str(PIP), "install", *need],
+                stdout=subprocess.DEVNULL, stderr=subprocess.PIPE,
+            )
+        except subprocess.CalledProcessError as e:
+            emit_error(f"pip install failed: {e}")
+            sys.exit(1)
+        # Force reimport
+        for mod in ("numpy", "torch", "soundfile", "transformers"):
+            if mod in sys.modules:
+                del sys.modules[mod]
+
+    # NeMo toolkit is large and optional — try to install it but fall
+    # back gracefully if it's unavailable on this platform.
+    try:
+        import nemo.collections.asr  # noqa: F401
+    except ImportError:
+        emit_status("Installing nemo_toolkit[asr] (large — may take a few minutes)...")
+        try:
+            subprocess.check_call(
+                [str(PIP), "install", "nemo_toolkit[asr]"],
+                stdout=subprocess.DEVNULL, stderr=subprocess.PIPE,
+                timeout=600,
+            )
+        except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:
+            emit_status(f"NeMo install skipped ({e}) — will use transformers fallback")
+
+# ---------------------------------------------------------------------------
+# Bootstrap: re-exec inside venv
+# ---------------------------------------------------------------------------
+
+# --check short-circuit — runs on the host Python without any venv or
+# dependency install so CI and smoke tests can verify the script parses
+# + is callable without triggering a 5-minute NeMo download.
+if "--check" in sys.argv:
+    emit({"type": "check", "ok": True, "script": str(Path(__file__).resolve())})
+    sys.exit(0)
+
+if not _in_venv():
+    _ensure_venv()
+    os.execv(str(PY), [str(PY)] + sys.argv)
+
+_ensure_deps()
+
+# Now safe to import
+import numpy as np  # noqa: E402
+
+# ---------------------------------------------------------------------------
+# Backend loaders
+# ---------------------------------------------------------------------------
+
+def _load_nemo_model(model_id: str = MODEL_ID, force_cpu: bool = False):
+    """Try to load via NeMo toolkit. Returns (model, device) or (None, None).
+
+    Handles the "cuDNN not compatible with SM < 7.5" error by retrying
+    on CPU. This is the common failure mode on older NVIDIA GPUs where
+    the installed torch has a newer cuDNN than the hardware supports.
+    """
+    # If caller asked for CPU explicitly, hide the GPU from torch before
+    # importing anything that might touch CUDA.
+    if force_cpu:
+        os.environ["CUDA_VISIBLE_DEVICES"] = ""
+    try:
+        import nemo.collections.asr as nemo_asr
+        import torch
+    except ImportError:
+        return (None, None)
+    try:
+        emit_status(f"Loading NeMo model {model_id}...")
+        model = nemo_asr.models.ASRModel.from_pretrained(model_id)
+        model.eval()
+        # Force CPU to avoid cuDNN version mismatches on older GPUs
+        if force_cpu or not torch.cuda.is_available():
+            try:
+                model = model.cpu()
+            except Exception:
+                pass
+        return (model, "cpu" if (force_cpu or not torch.cuda.is_available()) else "cuda")
+    except Exception as e:
+        msg = str(e)
+        emit_status(f"NeMo load failed: {msg[:200]}")
+        # Retry on CPU if the error looks like a cuDNN / device compat issue
+        if not force_cpu and any(k in msg for k in ("cuDNN", "SM <", "CUDA", "device side")):
+            emit_status("Retrying NeMo load on CPU only...")
+            return _load_nemo_model(model_id, force_cpu=True)
+        return (None, None)
+
+
+def _load_transformers_model(model_id: str = MODEL_ID):
+    """Fallback: load via HuggingFace transformers pipeline."""
+    try:
+        from transformers import pipeline
+    except ImportError:
+        return None
+    try:
+        emit_status(f"Loading transformers pipeline for {model_id}...")
+        device = -1
+        try:
+            import torch
+            if torch.cuda.is_available():
+                device = 0
+        except ImportError:
+            pass
+        pipe = pipeline(
+            task="automatic-speech-recognition",
+            model=model_id,
+            device=device,
+            return_timestamps=False,
+            chunk_length_s=30,
+            stride_length_s=5,
+        )
+        return pipe
+    except Exception as e:
+        emit_status(f"transformers load failed: {e}")
+        return None
+
+
+def _extract_hypothesis_text(r0) -> str:
+    """Extract the transcript string from a NeMo result item. Handles
+    plain strings, Hypothesis objects (with possibly empty text), and
+    nested lists of Hypotheses returned by RNNT models. Returns an
+    empty string for silent input rather than dumping the repr."""
+    if r0 is None:
+        return ""
+    if isinstance(r0, str):
+        return r0.strip()
+    # Nested list of hypotheses (some RNNT decoders)
+    if isinstance(r0, list):
+        if not r0:
+            return ""
+        return _extract_hypothesis_text(r0[0])
+    # Hypothesis object — may have text="" for silent audio, which is
+    # a VALID transcript (just empty). Return it without falling through
+    # to str(r0) which would dump the whole repr.
+    if hasattr(r0, "text"):
+        return str(r0.text or "").strip()
+    # best_hypothesis() method (rare)
+    if hasattr(r0, "best_hypothesis"):
+        try:
+            bh = r0.best_hypothesis()
+            if bh and hasattr(bh, "text"):
+                return str(bh.text or "").strip()
+        except Exception:
+            pass
+    return ""
+
+
+def _transcribe_buffer_nemo(model, audio: np.ndarray) -> str:
+    """Transcribe a 16kHz mono float32 numpy array via NeMo.
+
+    Tries multiple invocation signatures across NeMo versions:
+      - transcribe([np.ndarray]) — newest
+      - transcribe(paths2audio_files=["file.wav"]) — legacy, requires tmp wav
+    """
+    try:
+        # Newest API: pass audio arrays directly
+        result = model.transcribe([audio], batch_size=1, verbose=False)
+        if not result:
+            return ""
+        return _extract_hypothesis_text(result[0])
+    except TypeError:
+        # Older NeMo — fall through to file path invocation
+        pass
+    except Exception as e:
+        emit_error(f"NeMo transcribe error: {e}")
+        return ""
+
+    # Fallback: write audio to a temp WAV and pass the path
+    try:
+        import soundfile as sf
+        import tempfile
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
+            sf.write(tmp.name, audio, SAMPLE_RATE, subtype="PCM_16")
+            tmp_path = tmp.name
+        try:
+            result = model.transcribe(paths2audio_files=[tmp_path], batch_size=1, verbose=False)
+            if result and result[0] is not None:
+                return _extract_hypothesis_text(result[0])
+        finally:
+            try:
+                os.unlink(tmp_path)
+            except Exception:
+                pass
+    except Exception as e:
+        emit_error(f"NeMo legacy transcribe error: {e}")
+    return ""
+
+
+def _transcribe_buffer_transformers(pipe, audio: np.ndarray) -> str:
+    """Transcribe via transformers pipeline."""
+    try:
+        result = pipe({"array": audio, "sampling_rate": SAMPLE_RATE})
+        if isinstance(result, dict):
+            return str(result.get("text", "")).strip()
+        if isinstance(result, list) and result:
+            return str(result[0].get("text", "")).strip() if isinstance(result[0], dict) else ""
+        return ""
+    except Exception as e:
+        emit_error(f"transformers transcribe error: {e}")
+        return ""
+
+# ---------------------------------------------------------------------------
+# File transcription mode (single-shot)
+# ---------------------------------------------------------------------------
+
+def transcribe_file(path: str, language: str = "en") -> int:
+    """Single-file transcription — reads a WAV, prints one transcript
+    JSON line, exits. Used by AsrListenTool's file path. Exit code 0
+    on success, 1 on failure."""
+    try:
+        import soundfile as sf
+        audio, sr = sf.read(path, dtype="float32")
+        if audio.ndim > 1:
+            audio = audio.mean(axis=1)  # downmix to mono
+        if sr != SAMPLE_RATE:
+            # Resample via simple linear interpolation (avoids scipy dep)
+            ratio = SAMPLE_RATE / sr
+            new_len = int(len(audio) * ratio)
+            idx = np.linspace(0, len(audio) - 1, new_len).astype(np.float32)
+            audio = np.interp(idx, np.arange(len(audio), dtype=np.float32), audio).astype(np.float32)
+    except Exception as e:
+        emit_error(f"Failed to load audio file {path}: {e}")
+        return 1
+
+    (model, device) = _load_nemo_model()
+    backend = "nemo"
+    if model is None:
+        model = _load_transformers_model()
+        backend = "transformers"
+    if model is None:
+        emit_error("No nemotron backend available (tried NeMo + transformers)")
+        return 1
+
+    emit({"type": "ready", "backend": backend, "device": device or "cpu"})
+
+    t0 = time.time()
+    if backend == "nemo":
+        text = _transcribe_buffer_nemo(model, audio)
+    else:
+        text = _transcribe_buffer_transformers(model, audio)
+    elapsed = time.time() - t0
+
+    # Silent / no-speech audio is NOT an error — it's a valid transcript
+    # (empty string). Emit the full envelope so the caller can distinguish
+    # "no speech" from "engine crashed". Exit 0 either way.
+    emit({
+        "type": "transcript",
+        "text": text or "",
+        "isFinal": True,
+        "backend": f"nemotron-{backend}",
+        "latencyMs": int(elapsed * 1000),
+        "audioSeconds": float(len(audio) / SAMPLE_RATE),
+        "empty": not bool(text),
+    })
+    return 0
+
+# ---------------------------------------------------------------------------
+# Streaming mode (stdin → transcripts)
+# ---------------------------------------------------------------------------
+
+def stream_stdin(args) -> int:
+    (model, _device) = _load_nemo_model(args.model)
+    backend = "nemo"
+    if model is None:
+        model = _load_transformers_model(args.model)
+        backend = "transformers"
+    if model is None:
+        emit_error("No nemotron backend available (tried NeMo + transformers)")
+        return 1
+
+    emit({"type": "ready"})
+
+    audio_buf = np.zeros(0, dtype=np.float32)
+    buf_lock = threading.Lock()
+    chunk_bytes = int(args.chunk_seconds * SAMPLE_RATE * SAMPLE_WIDTH)
+    window_samples = int(args.window_seconds * SAMPLE_RATE)
+    last_text = ""
+    running = True
+
+    def read_stdin():
+        nonlocal audio_buf, running
+        try:
+            while running:
+                data = sys.stdin.buffer.read(chunk_bytes)
+                if not data:
+                    break
+                samples = np.frombuffer(data, dtype=np.int16).astype(np.float32) / 32768.0
+                with buf_lock:
+                    audio_buf = np.concatenate([audio_buf, samples])
+        except Exception:
+            pass
+        finally:
+            running = False
+
+    reader = threading.Thread(target=read_stdin, daemon=True)
+    reader.start()
+
+    try:
+        while running:
+            time.sleep(args.chunk_seconds)
+            with buf_lock:
+                if len(audio_buf) < SAMPLE_RATE:
+                    continue
+                window = audio_buf[-window_samples:].copy() if len(audio_buf) > window_samples else audio_buf.copy()
+            if backend == "nemo":
+                text = _transcribe_buffer_nemo(model, window)
+            else:
+                text = _transcribe_buffer_transformers(model, window)
+            if text and text != last_text:
+                last_text = text
+                emit_transcript(text, is_final=False, backend=f"nemotron-{backend}")
+    except KeyboardInterrupt:
+        pass
+
+    with buf_lock:
+        full_audio = audio_buf.copy()
+    if len(full_audio) >= SAMPLE_RATE:
+        if backend == "nemo":
+            text = _transcribe_buffer_nemo(model, full_audio)
+        else:
+            text = _transcribe_buffer_transformers(model, full_audio)
+        if text:
+            emit_transcript(text, is_final=True, backend=f"nemotron-{backend}")
+    running = False
+    return 0
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def main():
+    parser = argparse.ArgumentParser(description="Nemotron streaming ASR worker")
+    parser.add_argument("--model", default=MODEL_ID, help="HuggingFace model id (default: nvidia/nemotron-speech-streaming-en-0.6b)")
+    parser.add_argument("--file", default=None, help="Transcribe a single audio file instead of stdin")
+    parser.add_argument("--language", default="en", help="Language code")
+    parser.add_argument("--chunk-seconds", type=float, default=CHUNK_SECONDS, help="Transcribe interval")
+    parser.add_argument("--window-seconds", type=float, default=WINDOW_SECONDS, help="Sliding window size")
+    parser.add_argument("--stdin", action="store_true", help="Explicit stdin mode (default when no --file)")
+    parser.add_argument("--check", action="store_true", help="Just verify the script parses + imports; no model load")
+    args = parser.parse_args()
+
+    if args.check:
+        emit({"type": "check", "ok": True, "script": str(Path(__file__).resolve())})
+        return 0
+
+    if args.file:
+        return transcribe_file(args.file, args.language)
+    return stream_stdin(args)
+
+
+if __name__ == "__main__":
+    sys.exit(main())