stream_mn.py

#!/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()