Spaces:
Running on Zero
Running on Zero
| import os | |
| os.environ.setdefault("NUMBA_DISABLE_CUDA", "1") | |
| import spaces # noqa: E402 (must precede torch) | |
| import json | |
| import subprocess | |
| import sys | |
| import sysconfig | |
| import time | |
| from pathlib import Path | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| from huggingface_hub import snapshot_download | |
| MODEL_ID = "LocalDoc/speakwrite-azerbaijani" | |
| HERE = Path(__file__).resolve().parent | |
| BALBAL = HERE / "balbal" # balbal @ 1be9f88b77c63fbafa7713c6e46e79847f6c390e (Apache-2.0), refiner.json from the model repo | |
| NATIVE = BALBAL / "inference" / "native" | |
| MAX_SECONDS = 300 | |
| CHUNK_SECONDS = 30.0 | |
| # ---- build balbal's native C++/pybind11 KenLM decoders once (flashlight-text 0.0.7 ABI) ---- | |
| if not (NATIVE / ("ctc_fast_decoder" + sysconfig.get_config_var("EXT_SUFFIX"))).exists(): | |
| print("Building balbal native decoders ...", flush=True) | |
| subprocess.run([sys.executable, str(BALBAL / "inference" / "build_native.py")], check=True, | |
| env=dict(os.environ, CXXFLAGS="-O3")) | |
| sys.path[:0] = [str(BALBAL / "inference"), str(BALBAL)] | |
| from balbal import models as M # noqa: E402 | |
| from balbal.audio import read_audio # noqa: E402 | |
| from decoder import KenLMDecoder # noqa: E402 | |
| # ---- weights ---- | |
| model_dir = Path(snapshot_download(MODEL_ID)) | |
| settings = dict(M.decoder_config("balbal-v2-manas"), **json.loads((model_dir / "decoder.json").read_text())) | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| net = M.build_model("balbal-v2-manas", model_dir / "tokenizer", model_dir / "model.pt", device="cpu") | |
| net = net.to(DEVICE).eval() | |
| torch.set_float32_matmul_precision("highest") | |
| lm_decoder = KenLMDecoder(model_dir / "lm", settings, fast=True) | |
| sp = lm_decoder.sp | |
| VOCAB = lm_decoder.vocab_size | |
| print(f"Model + KenLM loaded (device={DEVICE}).", flush=True) | |
| def _forward(wave: np.ndarray) -> np.ndarray: | |
| """Mono float32 16 kHz waveform (<= 30 s) -> [frames, V + 1] natural-log CTC probabilities (blank last).""" | |
| wave = np.ascontiguousarray(wave, dtype=np.float32) | |
| if len(wave) < 4000: # the conv subsampler needs >= 0.25 s | |
| wave = np.pad(wave, (0, 4000 - len(wave))) | |
| with torch.inference_mode(): | |
| lp, n, _ = net(input_signal=torch.from_numpy(wave)[None].to(DEVICE), | |
| input_signal_length=torch.tensor([len(wave)], device=DEVICE)) | |
| return lp[0, : int(n[0])].float().cpu().numpy() | |
| def _greedy(lp: np.ndarray) -> str: | |
| ids = lp.argmax(axis=1) | |
| out, prev = [], -1 | |
| for t in ids: | |
| t = int(t) | |
| if t != prev and t < VOCAB: | |
| out.append(t) | |
| prev = t | |
| return sp.decode(out) | |
| def transcribe(audio_path: str, use_lm: bool = True, hotwords: str = ""): | |
| """Transcribe Azerbaijani speech to text with speakwrite-azerbaijani (FastConformer-CTC + 6-gram KenLM). | |
| Args: | |
| audio_path: path to an audio file (any format ffmpeg can read); resampled to 16 kHz mono. Max 5 minutes. | |
| use_lm: decode with the 6-gram token KenLM beam search (lower WER). If False, greedy CTC argmax. | |
| hotwords: optional comma-separated names/words expected in the audio; they get a small bonus during LM decoding. | |
| Returns: | |
| The transcript (cased, punctuated, numbers as words), word timestamps, and timing info. | |
| """ | |
| if not audio_path: | |
| raise gr.Error("Please upload or record an audio clip.") | |
| wave = read_audio(audio_path) | |
| duration = len(wave) / 16000 | |
| if duration > MAX_SECONDS: | |
| raise gr.Error(f"Audio is {duration:.0f} s long; please keep it under {MAX_SECONDS // 60} minutes.") | |
| hot = [w.strip() for w in (hotwords or "").split(",") if w.strip()] or None | |
| t0 = time.perf_counter() | |
| step = int(CHUNK_SECONDS * 16000) | |
| texts, words = [], [] | |
| for i in range(0, len(wave), step): | |
| lp = _forward(wave[i : i + step]) | |
| if use_lm: | |
| out = lm_decoder.decode_words(lp, hot, i / 16000) | |
| words += [[w["word"], round(w["start"], 2), round(w["end"], 2)] for w in out["words"]] | |
| text = out["text"] | |
| else: | |
| text = _greedy(lp) | |
| if text: | |
| texts.append(text) | |
| elapsed = time.perf_counter() - t0 | |
| info = (f"Audio: {duration:.1f} s · inference: {elapsed:.2f} s " | |
| f"({duration / elapsed if elapsed else 0:.0f}× real time) · decoder: {'KenLM beam search' if use_lm else 'greedy CTC'}") | |
| return " ".join(texts), words, info | |
| 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( | |
| "# 🎙️ speakwrite-azerbaijani\n" | |
| "Azerbaijani speech recognition: FastConformer-CTC with a 3-pass loop refiner, product-key acoustic memory " | |
| "and a 6-gram token KenLM. Output is cased and punctuated, with numbers written as words. \n" | |
| "[Model card](https://huggingface.co/LocalDoc/speakwrite-azerbaijani) · " | |
| "[balbal](https://github.com/erayyap/balbal) · " | |
| "*speakwrite-azerbaijani by LocalDoc* — CC BY-NC 4.0 (non-commercial use only)." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| audio = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Azerbaijani speech") | |
| run = gr.Button("Transcribe", variant="primary") | |
| with gr.Accordion("Advanced settings", open=False): | |
| use_lm = gr.Checkbox(value=True, label="Use KenLM beam search (uncheck for greedy CTC)") | |
| hotwords = gr.Textbox(value="", label="Hotwords (comma-separated names/words to favour; LM only)", | |
| placeholder="e.g. Bakı, Nizami") | |
| with gr.Column(): | |
| text = gr.Textbox(label="Transcript", lines=6, buttons=["copy"]) | |
| info = gr.Markdown() | |
| words = gr.Dataframe(headers=["word", "start (s)", "end (s)"], label="Word timestamps", interactive=False) | |
| run.click(transcribe, inputs=[audio, use_lm, hotwords], outputs=[text, words, info]) | |
| gr.Examples( | |
| examples=[["examples/suni_zeka.wav"], ["examples/honkonq.wav"], | |
| ["examples/medeni_sok.wav"], ["examples/quslar.wav"]], | |
| inputs=[audio], | |
| outputs=[text, words, info], | |
| fn=transcribe, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| label="Examples (FLEURS Azerbaijani test set, CC BY 4.0)", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |