#!/usr/bin/env python3 """ stream_mn.py - Mongolian Moonshine streaming inference. Modes: File: python stream_mn.py --model ./results-moonshine-mn/final --audio test.wav Live: python stream_mn.py --model ./results-moonshine-mn/final --live HF Hub: python stream_mn.py --model orgilj/moonshine-mn --audio test.wav Streaming works by VAD-segmenting the microphone input and transcribing each speech segment independently using MoonshineForConditionalGeneration.generate(). """ import argparse, sys, time from pathlib import Path import numpy as np, torch, soundfile as sf from transformers import MoonshineForConditionalGeneration, AutoFeatureExtractor sys.path.insert(0, str(Path(__file__).parent.parent)) from moonshine_ft.mn_tokenizer import MnBPETokenizer SR = 16000 TARGET_RMS = 0.075 BOS_ID, EOS_ID, PAD_ID = 1, 2, 2 def rms_norm(wav): rms = np.sqrt(np.mean(wav ** 2)) + 1e-8 return wav * (TARGET_RMS / rms) def load_audio(path): wav, sr = sf.read(str(path), dtype="float32", always_2d=False) if wav.ndim > 1: wav = wav.mean(1) if sr != SR: import librosa wav = librosa.resample(wav, orig_sr=sr, target_sr=SR) return np.ascontiguousarray(wav, dtype=np.float32) class MnASR: def __init__(self, model_path, bpe_path=None, device=None): self.dev = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu")) print(f"Loading model from: {model_path} [{self.dev}]") self.model = MoonshineForConditionalGeneration.from_pretrained(model_path).to(self.dev).eval() self.fe = AutoFeatureExtractor.from_pretrained( model_path if Path(model_path).is_dir() else "UsefulSensors/moonshine-base" ) if bpe_path and Path(bpe_path).exists(): self.tok = MnBPETokenizer(vocab_file=bpe_path) else: # Try loading from model dir or HF Hub try: self.tok = MnBPETokenizer.from_pretrained(model_path) except Exception: raise RuntimeError( "Cannot load tokenizer. Pass --bpe /path/to/mn_bpe.model " "or use a HF Hub model that includes mn_bpe.model." ) gc = self.model.generation_config gc.bos_token_id = BOS_ID gc.eos_token_id = EOS_ID gc.pad_token_id = PAD_ID gc.decoder_start_token_id = BOS_ID gc.max_length = None print(f"Model ready. Vocab size: {self.tok.vocab_size}") @torch.inference_mode() def transcribe(self, wav, max_new_tokens=None): wav = rms_norm(wav) dur = len(wav) / SR if max_new_tokens is None: max_new_tokens = max(5, min(int(dur * 8), 120)) inp = self.fe(wav, sampling_rate=SR, return_tensors="pt") iv = inp.input_values.to(self.dev) ids = self.model.generate(iv, max_new_tokens=max_new_tokens) return self.tok.decode_ids(ids[0].tolist()) def transcribe_file(model, path): wav = load_audio(path) dur = len(wav) / SR t0 = time.time() text = model.transcribe(wav) elapsed = time.time() - t0 rtf = elapsed / dur print(f"[{dur:.1f}s audio | {elapsed:.2f}s decode | RTF {rtf:.2f}x]") print(f"Transcription: {text}") return text def live_stream(model, chunk_sec=3.0, use_vad=True): try: import sounddevice as sd except ImportError: sys.exit("Install sounddevice: pip install sounddevice") vad_model, vad_iter = None, None if use_vad: try: vad_model, utils = torch.hub.load( "snakers4/silero-vad", "silero_vad", source="github", onnx=True ) (_, _, _, VADIterator, _) = utils vad_iter = VADIterator(vad_model) print("VAD loaded (Silero)") except Exception as e: print(f"VAD unavailable ({e}), using fixed-chunk mode") buf = np.array([], dtype=np.float32) speech_buf = np.array([], dtype=np.float32) is_speaking = False vad_buf = np.array([], dtype=np.float32) def callback(indata, frames, time_info, status): nonlocal buf, speech_buf, is_speaking, vad_buf chunk = indata[:, 0].astype(np.float32) if vad_iter is not None: if is_speaking: speech_buf = np.append(speech_buf, chunk) vad_buf = np.append(vad_buf, chunk) if len(vad_buf) >= 1536 * 3: ev = vad_iter(vad_buf, return_seconds=True) vad_buf = np.array([], dtype=np.float32) if ev: if "start" in ev: speech_buf = np.array([], dtype=np.float32) is_speaking = True print(" [speech start]") elif "end" in ev: is_speaking = False if len(speech_buf) > SR * 0.3: text = model.transcribe(speech_buf) if text: print(f">> {text}") speech_buf = np.array([], dtype=np.float32) else: buf = np.append(buf, chunk) if len(buf) >= int(SR * chunk_sec): text = model.transcribe(buf) if text: print(f">> {text}") buf = np.array([], dtype=np.float32) print("\nLive transcription started. Speak into microphone. Ctrl-C to stop.\n") with sd.InputStream(samplerate=SR, channels=1, callback=callback): try: while True: sd.sleep(500) except KeyboardInterrupt: print("\nStopped.") def main(): ap = argparse.ArgumentParser(description="Mongolian Moonshine ASR") ap.add_argument("--model", required=True, help="Local checkpoint dir or HF Hub id (orgilj/moonshine-mn)") ap.add_argument("--bpe", default="/workspace/mn_bpe.model", help="Path to mn_bpe.model (only needed for local checkpoint)") ap.add_argument("--audio", help="Audio file to transcribe") ap.add_argument("--live", action="store_true", help="Live mic transcription") ap.add_argument("--no-vad", action="store_true", help="Disable VAD in live mode") ap.add_argument("--chunk", type=float, default=3.0, help="Chunk seconds for fixed-chunk live mode") ap.add_argument("--device", choices=["cuda", "cpu"], help="Force device") a = ap.parse_args() if not a.audio and not a.live: ap.error("Specify --audio FILE or --live") asr = MnASR(a.model, bpe_path=a.bpe, device=a.device) if a.live: live_stream(asr, chunk_sec=a.chunk, use_vad=not a.no_vad) else: for audio_path in a.audio.split(","): p = Path(audio_path.strip()) if p.is_file(): transcribe_file(asr, p) elif p.is_dir(): files = sorted(p.glob("*.wav")) + sorted(p.glob("*.flac")) for f in files: print(f"\n--- {f.name} ---") transcribe_file(asr, f) else: print(f"Not found: {p}") if __name__ == "__main__": main()