Spaces:
Running on Zero
Running on Zero
| #!/usr/bin/env python3 | |
| """ | |
| run_inference.py: Production-ready Korean STT Inference Engine for Notesoop-STT-Ko-v1. | |
| - Google EmbeddingGemma-2 Frozen 300M Audio Tower (MRL 256-dim feature extraction) | |
| - 10-Layer Conformer Reader (21.5M params) on Apple Silicon MPS / CUDA / CPU | |
| - 25,000 Morpheme Vocabulary with Cosine Similarity Matching | |
| - Asynchronous KenLM 3-Gram Language Model Beam Search Rescoring | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import difflib | |
| import os | |
| import re | |
| import subprocess | |
| import sys | |
| import time | |
| from pathlib import Path | |
| from typing import Optional, List, Tuple | |
| from concurrent.futures import ThreadPoolExecutor, Future | |
| _ROOT = Path(__file__).resolve().parent | |
| _SRC_DIR = _ROOT / "src" | |
| if str(_SRC_DIR) not in sys.path: | |
| sys.path.insert(0, str(_SRC_DIR)) | |
| if str(_ROOT) not in sys.path: | |
| sys.path.insert(0, str(_ROOT)) | |
| import numpy as np | |
| import torch | |
| import yaml | |
| from dataset import KoreanMorphemeTokenizer, normalize_text | |
| from model import EmbeddingGemmaCTCReader | |
| from step2_extract_features import GenuineAudioFeatureExtractor | |
| from decoder import KenLMBeamSearchDecoder | |
| def load_audio( | |
| file_path: str, | |
| target_sr: int = 16000, | |
| start_sec: float = 0.0, | |
| duration_limit: float = 0.0 | |
| ) -> Tuple[torch.Tensor, float]: | |
| """ | |
| Loads any audio container (.wav, .m4a, .mp3, .flac) into 16kHz float32 mono tensor via ffmpeg. | |
| Uses ultra-fast seeking (-ss / -t) directly in FFmpeg to eliminate reading massive containers into memory. | |
| """ | |
| p = Path(file_path) | |
| if not p.exists(): | |
| raise FileNotFoundError(f"Audio file not found: {file_path}") | |
| cmd = ["ffmpeg", "-nostdin", "-threads", "0"] | |
| if start_sec > 0: | |
| cmd.extend(["-ss", str(start_sec)]) | |
| if duration_limit > 0: | |
| cmd.extend(["-t", str(duration_limit)]) | |
| cmd.extend([ | |
| "-i", str(p), | |
| "-f", "f32le", "-ac", "1", "-ar", str(target_sr), "-" | |
| ]) | |
| try: | |
| proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL) | |
| out, _ = proc.communicate() | |
| if proc.returncode != 0 or len(out) == 0: | |
| raise RuntimeError("FFmpeg decode returned empty data") | |
| arr = np.frombuffer(out, dtype=np.float32).copy() | |
| except Exception: | |
| # Fallback to soundfile if ffmpeg is not installed | |
| import soundfile as sf | |
| arr, sr = sf.read(str(p), dtype="float32") | |
| if arr.ndim > 1: | |
| arr = arr.mean(axis=1) | |
| if start_sec > 0: | |
| arr = arr[int(start_sec * sr):] | |
| if duration_limit > 0: | |
| arr = arr[:int(duration_limit * sr)] | |
| if sr != target_sr: | |
| import math | |
| gcd = math.gcd(sr, target_sr) | |
| up = target_sr // gcd | |
| down = sr // gcd | |
| import scipy.signal as sps | |
| arr = sps.resample_poly(arr, up, down).astype(np.float32) | |
| wav = torch.from_numpy(arr).unsqueeze(0) | |
| duration_sec = wav.size(-1) / float(target_sr) | |
| return wav, duration_sec | |
| def merge_chunk_transcripts(chunks: List[str]) -> str: | |
| """ | |
| Merges overlapping transcript segments by finding maximum character/word overlap. | |
| Prevents stuttering/duplication across sliding audio window boundaries. | |
| """ | |
| if not chunks: | |
| return "" | |
| if len(chunks) == 1: | |
| return chunks[0] | |
| merged = chunks[0] | |
| for nxt in chunks[1:]: | |
| if not nxt: | |
| continue | |
| max_search = min(len(merged), len(nxt), 40) | |
| found_overlap = 0 | |
| # Exact substring suffix-prefix match | |
| for k in range(max_search, 2, -1): | |
| if merged.endswith(nxt[:k]): | |
| found_overlap = k | |
| break | |
| # Fuzzy sequence matcher | |
| if found_overlap == 0 and len(merged) >= 10 and len(nxt) >= 10: | |
| tail = merged[-30:] | |
| head = nxt[:30] | |
| matcher = difflib.SequenceMatcher(None, tail, head) | |
| match = matcher.find_longest_match(0, len(tail), 0, len(head)) | |
| if match.size >= 4 and match.a + match.size == len(tail) and match.b == 0: | |
| found_overlap = match.size | |
| if found_overlap > 0: | |
| merged = merged + nxt[found_overlap:] | |
| else: | |
| merged = merged + " " + nxt | |
| return re.sub(r"\s+", " ", merged).strip() | |
| class EmbeddingGemmaSTT: | |
| """ | |
| High-performance, standalone Korean STT Pipeline for Notesoop-STT-Ko-v1. | |
| """ | |
| def __init__( | |
| self, | |
| checkpoint_path: Optional[str] = None, | |
| config_path: Optional[str] = None, | |
| vocab_path: Optional[str] = None, | |
| vector_bank_path: Optional[str] = None, | |
| lm_path: Optional[str] = None, | |
| device: str = "auto", | |
| use_kenlm: bool = True, | |
| lm_alpha: float = 0.12, | |
| lm_beta: float = 0.40, | |
| **kwargs # Ignore legacy kwargs like engine | |
| ): | |
| self.root = _ROOT | |
| self.lm_alpha = lm_alpha | |
| self.lm_beta = lm_beta | |
| # 0. Determine Device | |
| if device == "auto": | |
| dev_str = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu") | |
| else: | |
| dev_str = "mps" if (device == "metal" and torch.backends.mps.is_available()) else device | |
| self.device = torch.device(dev_str) | |
| # 1. Resolve paths & Tokenizer | |
| v_p = Path(vocab_path) if vocab_path else self.root / "data/vocab/vocab_25k.json" | |
| if not v_p.exists(): | |
| v_p = self.root / "data/vocab/vocab.json" | |
| self.tokenizer = KoreanMorphemeTokenizer(str(v_p)) | |
| vb_p = Path(vector_bank_path) if vector_bank_path else self.root / "data/vocab/vector_bank_25k.pt" | |
| if not vb_p.exists(): | |
| vb_p = self.root / "data/vocab/vector_bank.pt" | |
| self.vector_bank = torch.load(vb_p, map_location="cpu", weights_only=True) | |
| # 2. Initialize Conformer Reader | |
| if checkpoint_path: | |
| ckpt_p = Path(checkpoint_path) | |
| else: | |
| p_inf = self.root / "checkpoints/best_reader_inference.pt" | |
| p_full = self.root / "checkpoints/best_reader.pt" | |
| ckpt_p = p_inf if p_inf.exists() else p_full | |
| if not ckpt_p.exists(): | |
| raise FileNotFoundError(f"Reader checkpoint not found at: {ckpt_p}") | |
| ckpt = torch.load(ckpt_p, map_location="cpu", weights_only=False) | |
| state_dict = ckpt.get("model_state_dict", ckpt) | |
| vb_model = state_dict.get("ctc_head.vector_bank", self.vector_bank) | |
| d_model = state_dict["reader.input_proj.weight"].shape[0] if "reader.input_proj.weight" in state_dict else 384 | |
| input_dim = state_dict["reader.input_proj.weight"].shape[1] if "reader.input_proj.weight" in state_dict else 256 | |
| detected_layers = len([k for k in state_dict.keys() if k.startswith("reader.layers.") and k.endswith(".final_norm.weight")]) | |
| num_layers = detected_layers if detected_layers > 0 else 10 | |
| dim_ffn = state_dict["reader.layers.0.ffn1.0.weight"].shape[0] if "reader.layers.0.ffn1.0.weight" in state_dict else 1536 | |
| nhead = d_model // 64 | |
| output_dim = state_dict["reader.output_proj.0.weight"].shape[0] if "reader.output_proj.0.weight" in state_dict else 256 | |
| self.reader = EmbeddingGemmaCTCReader( | |
| input_dim=input_dim, | |
| output_dim=output_dim, | |
| d_model=d_model, | |
| nhead=nhead, | |
| num_layers=num_layers, | |
| dim_feedforward=dim_ffn, | |
| dropout=0.0, | |
| init_temperature=0.05, | |
| vector_bank=vb_model, | |
| use_spec_augment=False, | |
| trainable_vector_bank=False | |
| ).to(self.device) | |
| self.reader.load_state_dict(state_dict, strict=False) | |
| self.reader.eval() | |
| # 3. Audio Feature Extractor (Google EmbeddingGemma-2 300M Audio Tower) | |
| model_name = "google/embeddinggemma-2" | |
| self.extractor = GenuineAudioFeatureExtractor(device=self.device, model_name=model_name) | |
| # 4. Optional KenLM Decoder | |
| self.decoder_kenlm = None | |
| if use_kenlm: | |
| lm_file = Path(lm_path) if lm_path else self.root / "data/vocab/korean_lm.bin" | |
| if lm_file.exists(): | |
| self.decoder_kenlm = KenLMBeamSearchDecoder( | |
| tokenizer=self.tokenizer, | |
| lm_path=str(lm_file), | |
| beam_size=8, | |
| alpha=lm_alpha, | |
| beta=lm_beta | |
| ) | |
| def transcribe( | |
| self, | |
| audio_path: str, | |
| chunk_sec: float = 16.0, | |
| overlap_sec: float = 1.0, | |
| start_sec: float = 0.0, | |
| duration_limit: float = 0.0, | |
| batch_size: int = 4 | |
| ) -> Tuple[str, float]: | |
| """ | |
| Transcribes audio file of arbitrary length into Korean text. | |
| Optimized with Batched GPU Forward Passes and Asynchronous KenLM Rescoring: | |
| - CPU Mel Preprocessing | |
| - GPU 300M Audio Tower Forward (Batched) | |
| - GPU Conformer Reader + Top-32 Reduction (Batched) | |
| - Asynchronous KenLM 3-Gram Beam Search Rescoring | |
| Returns (transcript_text, rtf_speed). | |
| """ | |
| wav, eval_dur = load_audio( | |
| audio_path, | |
| target_sr=16000, | |
| start_sec=start_sec, | |
| duration_limit=duration_limit | |
| ) | |
| sr = 16000 | |
| step_sec = max(chunk_sec - overlap_sec, 0.5) if overlap_sec > 0 else chunk_sec | |
| step_samples = int(step_sec * sr) | |
| chunk_samples = int(chunk_sec * sr) | |
| total_samples = wav.size(-1) | |
| chunk_ranges: List[Tuple[int, int]] = [] | |
| c_start = 0 | |
| while c_start < total_samples: | |
| c_end = min(c_start + chunk_samples, total_samples) | |
| if (c_end - c_start) / sr >= 0.5: | |
| chunk_ranges.append((c_start, c_end)) | |
| if c_end == total_samples: | |
| break | |
| c_start += step_samples | |
| if not chunk_ranges: | |
| return "", 0.0 | |
| t0 = time.time() | |
| # 1. Preprocess Mel-spectrogram on CPU (pad last chunk to chunk_samples for uniform shapes) | |
| preprocessed_chunks = [] | |
| for cs, ce in chunk_ranges: | |
| c_wav = wav[:, cs:ce] | |
| if c_wav.size(-1) < chunk_samples: | |
| c_wav = torch.nn.functional.pad(c_wav, (0, chunk_samples - c_wav.size(-1))) | |
| preprocessed_chunks.append(self.extractor.preprocess_audio_cpu(c_wav)) | |
| # 2. Pipelined Batched Inference | |
| chunk_texts: List[str] = [] | |
| kenlm_futures: List[Future] = [] | |
| use_kenlm_async = (self.decoder_kenlm is not None) | |
| with ThreadPoolExecutor(max_workers=2) as executor: | |
| for b_idx in range(0, len(preprocessed_chunks), batch_size): | |
| b_items = preprocessed_chunks[b_idx : b_idx + batch_size] | |
| # Step A: Extract Audio Tower Features (GPU) | |
| if len(b_items) > 1 and len(set(it[3] for it in b_items)) == 1: | |
| feats = self.extractor.extract_batch_same_bucket(b_items) | |
| else: | |
| feats = [self.extractor.extract_from_preprocessed(*it) for it in b_items] | |
| # Step B: Conformer Reader Forward (Batched GPU) | |
| if len(feats) > 1 and len(set(f.size(0) for f in feats)) == 1: | |
| feats_batch = torch.stack(feats, dim=0).to(device=self.device, dtype=torch.float32) | |
| logits_batch = self.reader(feats_batch) | |
| else: | |
| logits_batch = [ | |
| self.reader(f.unsqueeze(0).to(device=self.device, dtype=torch.float32)).squeeze(0) | |
| for f in feats | |
| ] | |
| # Step C: Decoding (KenLM Beam Search or Greedy CTC) | |
| for b in range(len(feats)): | |
| logits = logits_batch[b] | |
| if use_kenlm_async: | |
| topk_lprobs, topk_ids = torch.topk(logits, k=32, dim=-1) | |
| blank_lprobs = logits[:, 0] | |
| fut = executor.submit( | |
| self.decoder_kenlm.decode_topk, | |
| topk_ids.detach().cpu().to(torch.int32).numpy(), | |
| topk_lprobs.detach().cpu().float().numpy(), | |
| blank_lprobs.detach().cpu().float().numpy(), | |
| alpha=self.lm_alpha, | |
| beta=self.lm_beta | |
| ) | |
| kenlm_futures.append(fut) | |
| else: | |
| pred_ids = torch.argmax(logits, dim=-1).cpu().tolist() | |
| txt = self.tokenizer.decode(pred_ids, remove_blank=True, collapse_repeats=True) | |
| chunk_texts.append(txt) | |
| if use_kenlm_async: | |
| chunk_texts = [fut.result() for fut in kenlm_futures] | |
| t_elapsed = time.time() - t0 | |
| rtf = t_elapsed / max(eval_dur, 0.001) | |
| valid_texts = [t.strip() for t in chunk_texts if t.strip()] | |
| if overlap_sec > 0: | |
| full_text = merge_chunk_transcripts(valid_texts) | |
| else: | |
| full_text = " ".join(valid_texts) | |
| return full_text, rtf | |
| def main(): | |
| parser = argparse.ArgumentParser(description="Notesoop-STT-Ko-v1 Standalone Inference") | |
| parser.add_argument("--audio", type=str, default="sample_audio.wav", help="Path to audio file (.wav, .m4a, .mp3, etc.)") | |
| parser.add_argument("--checkpoint", type=str, default=None, help="Path to checkpoint (.pt)") | |
| parser.add_argument("--device", type=str, default="auto", choices=["auto", "mps", "cuda", "cpu"], help="Inference device") | |
| parser.add_argument("--kenlm", action="store_true", default=True, help="Enable KenLM 3-gram language model rescoring") | |
| parser.add_argument("--no_kenlm", dest="kenlm", action="store_false", help="Disable KenLM (Greedy decoding only)") | |
| parser.add_argument("--duration", type=float, default=0.0, help="Duration limit in seconds (0 for full audio)") | |
| parser.add_argument("--chunk_sec", type=float, default=16.0, help="Chunk length in seconds") | |
| parser.add_argument("--batch_size", type=int, default=4, help="Batch size for parallel processing (default: 4)") | |
| parser.add_argument("--output", type=str, default=None, help="Optional output .txt file path") | |
| args = parser.parse_args() | |
| print("\n" + "=" * 70) | |
| print("🌲 Notesoop-STT-Ko-v1: Korean Speech Recognition Engine") | |
| print("=" * 70) | |
| print(f" • Audio File : {args.audio}") | |
| print(f" • Device : {args.device.upper()}") | |
| print(f" • Batch Size : {args.batch_size}") | |
| print(f" • KenLM LM : {'Enabled (8-beam rescoring)' if args.kenlm else 'Disabled (Greedy)'}") | |
| print("-" * 70) | |
| stt = EmbeddingGemmaSTT( | |
| checkpoint_path=args.checkpoint, | |
| device=args.device, | |
| use_kenlm=args.kenlm | |
| ) | |
| print("\n🚀 Transcribing...") | |
| text, rtf = stt.transcribe( | |
| args.audio, | |
| chunk_sec=args.chunk_sec, | |
| duration_limit=args.duration, | |
| batch_size=args.batch_size | |
| ) | |
| print("\n" + "=" * 70) | |
| print("📝 [Transcription Result]:") | |
| print("=" * 70) | |
| print(text) | |
| print("-" * 70) | |
| print(f"⏱️ Speed / RTF: {rtf:.4f} ({1.0 / max(rtf, 1e-4):.1f}x Real-time on {stt.device.type.upper()})") | |
| print("=" * 70 + "\n") | |
| if args.output: | |
| out_p = Path(args.output) | |
| out_p.parent.mkdir(parents=True, exist_ok=True) | |
| out_p.write_text(text, encoding="utf-8") | |
| print(f"💾 Saved transcription to: {out_p}\n") | |
| if __name__ == "__main__": | |
| main() | |