import spaces # must be imported before torch import os import sys import time import gradio as gr import torch import soundfile as sf # Full float32 math, matching the reference implementation (no TF32). torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False APP_DIR = os.path.dirname(os.path.abspath(__file__)) for _p in (APP_DIR, os.path.join(APP_DIR, "src")): if _p not in sys.path: sys.path.insert(0, _p) from huggingface_hub import hf_hub_download # noqa: E402 from run_inference import EmbeddingGemmaSTT # noqa: E402 MODEL_ID = "Phsoft-ai/Notesoop-STT-Ko-v1" MAX_SECONDS = 180.0 # --- Weights ----------------------------------------------------------------- # The reader checkpoint, morpheme vocab, vector bank and KenLM 3-gram LM all # live in the model repo; EmbeddingGemmaSTT takes explicit paths for each. _CKPT = hf_hub_download(MODEL_ID, "checkpoints/best_reader_inference.pt") _VOCAB = hf_hub_download(MODEL_ID, "data/vocab/vocab_25k.json") _VECTOR_BANK = hf_hub_download(MODEL_ID, "data/vocab/vector_bank_25k.pt") _KENLM = hf_hub_download(MODEL_ID, "data/vocab/korean_lm.bin") # --- Engine (loaded eagerly at module scope, onto the GPU) ------------------- STT = EmbeddingGemmaSTT( checkpoint_path=_CKPT, vocab_path=_VOCAB, vector_bank_path=_VECTOR_BANK, lm_path=_KENLM, device="cuda", use_kenlm=True, lm_alpha=0.12, lm_beta=0.40, ) def _estimate_duration(audio_path, *args, **kwargs): """Reserve GPU time from the input length (extras swallowed: Gradio injects progress=). Measured on this Space (zero-a10g), GPU-side compute plus the per-request worker fork: 5.4 s of audio -> 0.42 s, 6.9 s -> 1.27 s (coldest call), 44.7 s -> 1.18 s, and the 180 s input cap -> 5.46 s of compute (7.2 s wall). `2.0 s + 0.035 s per audio second` tracks that with a ~1.4x margin on the longest clip, so visitor quota is not wasted on short inputs. """ secs = MAX_SECONDS if isinstance(audio_path, str) and os.path.exists(audio_path): try: secs = sf.info(audio_path).duration except Exception: pass return max(4, min(8, int(round(2.0 + secs * 0.035)))) @spaces.GPU(duration=_estimate_duration) @torch.no_grad() def transcribe( audio_path: str, chunk_sec: float = 16.0, overlap_sec: float = 1.0, batch_size: int = 4, progress=gr.Progress(track_tqdm=True), ) -> tuple[str, str]: """Transcribe Korean speech to text with Notesoop-STT-Ko-v1. Runs the frozen EmbeddingGemma-2 audio encoder (MRL 256-d features), a 10-layer Conformer CTC reader over a 25k Korean morpheme vocabulary, and KenLM 3-gram beam-search rescoring. Args: audio_path: path to the audio file (wav/mp3/m4a/flac; resampled to 16 kHz mono). chunk_sec: sliding-window length in seconds (16 s is the reference setting). overlap_sec: overlap between consecutive windows, used to stitch them back together. batch_size: number of windows processed per batched forward pass. progress: Gradio progress reporter (injected automatically). Returns: A (transcript, stats) tuple: the Korean transcript and a short speed report. """ if not audio_path: raise gr.Error("Please upload or record some audio first.") info = sf.info(audio_path) if info.duration > MAX_SECONDS: raise gr.Error( f"Audio is {info.duration:.0f}s long; this demo accepts up to {int(MAX_SECONDS)}s." ) t0 = time.perf_counter() text, rtf = STT.transcribe( audio_path, chunk_sec=chunk_sec, overlap_sec=overlap_sec, batch_size=batch_size, ) elapsed = time.perf_counter() - t0 text = text.strip() or "(no speech recognised)" stats = ( f"{info.duration:.1f} s of audio · transcribed in {elapsed:.2f} s " f"({1.0 / max(rtf, 1e-4):.1f}× real-time) · KenLM 3-gram rescoring" ) return text, stats CSS = """ #col-container { max-width: 1000px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ with gr.Blocks() as demo: with gr.Column(elem_id="col-container"): gr.Markdown( "# Notesoop-STT-Ko-v1 — Korean Speech Recognition\n" "Ultra-fast Korean on-device STT from " "[Phsoft-ai/Notesoop-STT-Ko-v1](https://huggingface.co/Phsoft-ai/Notesoop-STT-Ko-v1): " "the **frozen EmbeddingGemma-2 audio tower** produces 256-dimensional MRL embeddings, a small " "**10-layer Conformer CTC reader** (21.5M params) maps them to a 25,000-token Korean morpheme " "vocabulary, and a **KenLM 3-gram** model rescoring the CTC beam. Korean only." ) with gr.Row(): with gr.Column(): audio = gr.Audio( label="Korean speech", type="filepath", sources=["upload", "microphone"], ) run = gr.Button("Transcribe", variant="primary") with gr.Column(): transcript = gr.Textbox(label="Transcript", lines=8, buttons=["copy"]) stats = gr.Markdown() with gr.Accordion("Advanced options", open=False): chunk_sec = gr.Slider( 4.0, 30.0, value=16.0, step=1.0, label="Chunk length (s)", info="Sliding-window length. 16 s is the reference setting.", ) overlap_sec = gr.Slider( 0.0, 4.0, value=1.0, step=0.5, label="Chunk overlap (s)", info="Overlap used to stitch consecutive windows together.", ) batch_size = gr.Slider( 1, 8, value=4, step=1, label="Batch size", info="Windows processed per batched forward pass.", ) gr.Examples( examples=[ ["examples/zeroth_01.wav"], ["examples/zeroth_02.wav"], ["examples/zeroth_03.wav"], ["examples/zeroth_04.wav"], ], inputs=[audio], outputs=[transcript, stats], fn=transcribe, cache_examples=True, cache_mode="lazy", ) gr.Markdown( "Example clips are Korean read-speech utterances from " "[Zeroth-Korean](https://huggingface.co/datasets/Bingsu/zeroth-korean) " "(CC BY 4.0), the benchmark this model reports on. " "Up to 180 s of audio per run." ) run.click( transcribe, inputs=[audio, chunk_sec, overlap_sec, batch_size], outputs=[transcript, stats], api_name="transcribe", ) if __name__ == "__main__": demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)