Spaces:
Running on Zero
Running on Zero
| 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)))) | |
| 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) | |