diff --git a/examples/whisper.youtube/parakeet b/examples/whisper.youtube/parakeet index cde3d9a7b..f65d083cd 100755 --- a/examples/whisper.youtube/parakeet +++ b/examples/whisper.youtube/parakeet @@ -5,4 +5,8 @@ # ./parakeet set -euo pipefail DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# numba cannot write its cache on the mounted filesystem under this dir +# (no inode-stable locator); redirect cache to /tmp. +export NUMBA_CACHE_DIR="${NUMBA_CACHE_DIR:-/tmp/numba-cache-$USER}" +mkdir -p "$NUMBA_CACHE_DIR" exec "$DIR/.venv-parakeet/bin/python" "$DIR/parakeet_transcribe.py" "$@" diff --git a/examples/whisper.youtube/parakeet_mic.py b/examples/whisper.youtube/parakeet_mic.py new file mode 100644 index 000000000..607f42d2d --- /dev/null +++ b/examples/whisper.youtube/parakeet_mic.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python3 +"""Live microphone streaming with Parakeet-tdt-0.6b-v3. + +Captures audio from the mic (pw-record or arecord), runs a sliding-window +transcription, and prints text progressively as new audio arrives. +""" + +import argparse +import os +import shutil +import signal +import subprocess +import sys +import tempfile +import time + +import numpy as np + +MODEL_NAME = "nvidia/parakeet-tdt-0.6b-v3" +SR = 16000 +DEFAULT_WINDOW_S = 8.0 +DEFAULT_SHIFT_S = 2.5 + + +def _norm(w: str) -> str: + return "".join(ch for ch in w.lower() if ch.isalnum()) + + +def _dedup_overlap(prev_text: str, curr_text: str) -> str: + prev_words = prev_text.split() + curr_words = curr_text.split() + prev_norm = [_norm(w) for w in prev_words] + curr_norm = [_norm(w) for w in curr_words] + max_n = min(len(prev_norm), len(curr_norm)) + for n in range(max_n, 0, -1): + if prev_norm[-n:] == curr_norm[:n]: + return " ".join(curr_words[n:]) + return curr_text + + +def _find_recorder(): + for p in ("pw-record", "/usr/bin/pw-record", "/usr/local/bin/pw-record"): + if shutil.which(p) if not os.path.isabs(p) else os.path.isfile(p): + return p + return shutil.which("arecord") + + +def _build_record_cmd(): + rec = _find_recorder() + if not rec: + raise RuntimeError("Neither pw-record nor arecord found") + if "pw-record" in rec: + return [rec, "--format", "s16", "--rate", str(SR), "--channels", "1", "-"] + return [rec, "-f", "S16_LE", "-r", str(SR), "-c", "1", "-t", "raw"] + + +def load_model(): + import nemo.collections.asr as nemo_asr + from nemo.utils import logging as nemo_logging + import torch + import logging + + nemo_logging.set_verbosity(nemo_logging.ERROR) + logging.getLogger("nemo_logger").setLevel(logging.ERROR) + print(f"Loading {MODEL_NAME} ...", flush=True) + asr = nemo_asr.models.ASRModel.from_pretrained(MODEL_NAME) + asr = asr.cuda() if torch.cuda.is_available() else asr + nemo_logging.set_verbosity(nemo_logging.ERROR) + return asr + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--window", type=float, default=DEFAULT_WINDOW_S, + help=f"sliding window length in seconds (default: {DEFAULT_WINDOW_S})") + parser.add_argument("--shift", type=float, default=DEFAULT_SHIFT_S, + help=f"window shift / step in seconds (default: {DEFAULT_SHIFT_S})") + args = parser.parse_args() + + if args.shift <= 0 or args.window <= 0: + sys.exit("window and shift must be positive") + if args.shift > args.window: + sys.exit("shift must be <= window") + + import soundfile as sf + + asr = load_model() + cmd = _build_record_cmd() + + win_samples = int(args.window * SR) + shift_samples = int(args.shift * SR) + shift_bytes = shift_samples * 2 # int16 mono + + tmp = tempfile.NamedTemporaryFile(suffix=".wav", prefix="parakeet-mic-", delete=False) + tmp_wav = tmp.name + tmp.close() + + print(f"Recorder: {cmd[0]}", flush=True) + print(f"Window: {args.window:.1f}s, shift: {args.shift:.1f}s " + f"(overlap: {args.window - args.shift:.1f}s)", flush=True) + print("--- speak into mic; Ctrl-C to stop ---\n", flush=True) + + proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL) + + def _shutdown(*_): + if proc.poll() is None: + proc.send_signal(signal.SIGINT) + try: + proc.wait(timeout=2) + except subprocess.TimeoutExpired: + proc.kill() + if os.path.exists(tmp_wav): + os.unlink(tmp_wav) + print("\n--- stopped ---", flush=True) + sys.exit(0) + + signal.signal(signal.SIGINT, _shutdown) + signal.signal(signal.SIGTERM, _shutdown) + + buffer = np.zeros(0, dtype=np.int16) + prev_text = "" + + try: + while True: + new_bytes = proc.stdout.read(shift_bytes) + if not new_bytes: + break + new_samples = np.frombuffer(new_bytes, dtype=np.int16) + buffer = np.concatenate([buffer, new_samples]) + if len(buffer) > win_samples: + buffer = buffer[-win_samples:] + if len(buffer) < shift_samples: + continue + + sf.write(tmp_wav, buffer, SR) + t0 = time.monotonic() + out = asr.transcribe([tmp_wav], timestamps=False, verbose=False) + elapsed = time.monotonic() - t0 + current = out[0].text.strip() + if not current: + continue + + new_part = _dedup_overlap(prev_text, current) if prev_text else current + if new_part.strip(): + print(new_part, end=" ", flush=True) + prev_text = current + finally: + _shutdown() + + +if __name__ == "__main__": + main() diff --git a/examples/whisper.youtube/parakeet_transcribe.py b/examples/whisper.youtube/parakeet_transcribe.py index 45c3a80dd..55b2a2255 100644 --- a/examples/whisper.youtube/parakeet_transcribe.py +++ b/examples/whisper.youtube/parakeet_transcribe.py @@ -1,14 +1,28 @@ #!/usr/bin/env python3 -"""Transcribe audio with NVIDIA Parakeet-tdt-0.6b-v3 (multilingual).""" +"""Transcribe audio with NVIDIA Parakeet-tdt-0.6b-v3 (multilingual). +Modes: + batch — transcribe the whole audio at once, print at the end (default). + stream — sliding window with LocalAgreement: progressive text output as + new audio arrives. Each window of WINDOW seconds is re-transcribed + every SHIFT seconds; tokens agreed between two consecutive windows + are "confirmed" and printed; the rest stays uncertain until the + next iteration confirms (or invalidates) it. +""" + +import argparse import os import subprocess import sys import tempfile -from pathlib import Path MODEL_NAME = "nvidia/parakeet-tdt-0.6b-v3" +# Sliding-window streaming defaults (in seconds). +# Matches whisper.cpp/examples/stream defaults. +DEFAULT_WINDOW_S = 10.0 +DEFAULT_SHIFT_S = 3.0 + def download_audio(url: str, out_path: str): subprocess.run([ @@ -24,46 +38,142 @@ def convert_wav(src: str, dst: str): ], check=True, capture_output=True) -def transcribe(wav_path: str, long_audio: bool = True) -> str: +def load_model(long_audio: bool = True): import nemo.collections.asr as nemo_asr + import torch print(f"Loading {MODEL_NAME} ...", flush=True) asr = nemo_asr.models.ASRModel.from_pretrained(MODEL_NAME) - asr = asr.cuda() if __import__("torch").cuda.is_available() else asr - + asr = asr.cuda() if torch.cuda.is_available() else asr if long_audio: asr.change_attention_model( self_attention_model="rel_pos_local_attn", att_context_size=[256, 256], ) + return asr + +def transcribe_batch(asr, wav_path: str) -> str: print(f"Transcribing {wav_path} ...", flush=True) out = asr.transcribe([wav_path], timestamps=False) return out[0].text -def main(): - if len(sys.argv) < 2: - print(f"Usage: {sys.argv[0]} ") - sys.exit(1) +def _norm(w: str) -> str: + """Normalize a word for dedup comparison: lower-case, strip punctuation.""" + return "".join(ch for ch in w.lower() if ch.isalnum()) - src = sys.argv[1] + +def _dedup_overlap(prev_text: str, curr_text: str) -> str: + """Return the new portion of curr_text by removing the trailing words of + prev_text that match the leading words of curr_text (overlap region). + Comparison is case- and punctuation-insensitive.""" + prev_words = prev_text.split() + curr_words = curr_text.split() + prev_norm = [_norm(w) for w in prev_words] + curr_norm = [_norm(w) for w in curr_words] + max_n = min(len(prev_norm), len(curr_norm)) + for n in range(max_n, 0, -1): + if prev_norm[-n:] == curr_norm[:n]: + return " ".join(curr_words[n:]) + return curr_text + + +def transcribe_sliding(asr, wav_path: str, window_s: float, shift_s: float) -> str: + """Sliding-window streaming: each window of WINDOW seconds is transcribed + every SHIFT seconds; words overlapping with the previous window are + de-duplicated and only the new tail is printed.""" + import soundfile as sf + + if shift_s <= 0 or window_s <= 0: + raise ValueError("window and shift must be positive") + if shift_s > window_s: + raise ValueError("shift must be <= window") + + audio, sr = sf.read(wav_path) + if sr != 16000: + raise ValueError(f"Expected 16kHz WAV, got {sr}Hz") + + n_samples = len(audio) + win = int(window_s * sr) + shift = int(shift_s * sr) + overlap_s = window_s - shift_s + + print( + f"Sliding window: {window_s:.1f}s window, {shift_s:.1f}s shift " + f"({overlap_s:.1f}s overlap), audio {n_samples / sr:.1f}s", + flush=True, + ) + print("--- streaming transcription ---", flush=True) + + tmp = tempfile.NamedTemporaryFile(suffix=".wav", prefix="parakeet-slide-", + delete=False) + tmp_wav = tmp.name + tmp.close() + + full_text_parts = [] + prev_text = "" + pos = 0 + + try: + while pos < n_samples: + end = min(pos + win, n_samples) + chunk = audio[pos:end] + sf.write(tmp_wav, chunk, sr) + + out = asr.transcribe([tmp_wav], timestamps=False, verbose=False) + current = out[0].text.strip() + + new_part = _dedup_overlap(prev_text, current) if prev_text else current + if new_part: + print(new_part, end=" ", flush=True) + full_text_parts.append(new_part) + + prev_text = current + pos += shift + + print(flush=True) + return " ".join(full_text_parts).strip() + finally: + if os.path.exists(tmp_wav): + os.unlink(tmp_wav) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("source", help="WAV file path or YouTube URL") + parser.add_argument("--stream", action="store_true", + help="streaming mode (sliding window + LocalAgreement)") + parser.add_argument("--window", type=float, default=DEFAULT_WINDOW_S, + help=f"sliding window length in seconds (default: {DEFAULT_WINDOW_S})") + parser.add_argument("--shift", type=float, default=DEFAULT_SHIFT_S, + help=f"window shift in seconds (default: {DEFAULT_SHIFT_S}); " + "overlap = window - shift") + args = parser.parse_args() + + src = args.source if src.startswith(("http://", "https://", "youtu", "www.")): - with tempfile.TemporaryDirectory(prefix="parakeet-") as tmp: - raw = os.path.join(tmp, "audio.%(ext)s") - wav = os.path.join(tmp, "audio_16k.wav") - print(f"Downloading {src} ...", flush=True) - download_audio(src, raw) - downloaded = [f for f in os.listdir(tmp) if f.startswith("audio.")] - print("Converting to 16kHz WAV ...", flush=True) - convert_wav(os.path.join(tmp, downloaded[0]), wav) - text = transcribe(wav) + tmp = tempfile.mkdtemp(prefix="parakeet-") + raw = os.path.join(tmp, "audio.%(ext)s") + wav = os.path.join(tmp, "audio_16k.wav") + print(f"Downloading {src} ...", flush=True) + download_audio(src, raw) + downloaded = [f for f in os.listdir(tmp) if f.startswith("audio.")] + print("Converting to 16kHz WAV ...", flush=True) + convert_wav(os.path.join(tmp, downloaded[0]), wav) + wav_path = wav else: - text = transcribe(src) + wav_path = src - print("\n=== Transcription ===") - print(text) + asr = load_model(long_audio=not args.stream) + + if args.stream: + text = transcribe_sliding(asr, wav_path, args.window, args.shift) + else: + text = transcribe_batch(asr, wav_path) + print("\n=== Transcription ===") + print(text) if __name__ == "__main__":