"""Japanese anime / game voice ASR with speaker labels and timestamps (MOSS-Transcribe-Diarize fine-tuned on galgame voices). python transcribe.py input.wav # prints [start-end] Sxx: text python transcribe.py voices\\ --out result.jsonl # a folder, one JSON line per file python transcribe.py ep01.wav --srt ep01.srt python transcribe.py input.wav --num_beams 4 --punct_bias 1.0 Decoding safeguards (all on by default, see README): * loop guard bans a token that would extend a 1-4 token cycle repeated 8 times (ぇぇぇ…, ーーー…) * length cap max_new_tokens = 20 per second of audio + 64 * temperature fallback re-decode at T=0.2..0.8 while a segment is too repetitive (gzip ratio > 1.8) * silence gate no output for audio quieter than -55 dBFS; segments over such audio are dropped Long input is cut at quiet points into windows of at most --window seconds (speaker labels restart per window). """ from __future__ import annotations import argparse import contextlib import copy import json import os import re import sys import tempfile import zlib from pathlib import Path import numpy as np import soundfile as sf import soxr import torch from transformers import AutoModelForCausalLM, AutoProcessor, LogitsProcessor, LogitsProcessorList HERE = Path(__file__).resolve().parent sys.path.insert(0, str(HERE)) # bundled inference_utils from inference_utils import build_transcription_messages, prepare_inputs # noqa: E402 SR = 16000 SEG = re.compile(r"\[(\d+(?:\.\d+)?)\]\[(S\d+)\](.*?)\[(\d+(?:\.\d+)?)\]", re.S) AUDIO_EXT = {".wav", ".flac", ".ogg", ".opus", ".mp3", ".m4a"} def load(path) -> np.ndarray: a, sr = sf.read(str(path), dtype="float32", always_2d=True) a = a.mean(axis=1) return soxr.resample(a, sr, SR) if sr != SR else a def parse(text: str) -> list[tuple[float, float, str, str]]: return [(float(s), float(e), spk, txt) for s, spk, txt, e in SEG.findall(text)] def frame_db(a: np.ndarray, hop: float = 0.05) -> np.ndarray: n = int(hop * SR) if len(a) < n: a = np.pad(a, (0, n - len(a))) f = a[: len(a) // n * n].reshape(-1, n) return 10 * np.log10((f ** 2).mean(axis=1) + 1e-12) def compression_ratio(text: str) -> float: parts = [s[3] for s in parse(text)] or [re.sub(r"\[[^\]]*\]", "", text)] def cr(t: str) -> float: b = t.encode("utf-8") return len(b) / max(len(zlib.compress(b)), 1) if b else 0.0 return max(cr(p) for p in parts) class LoopGuard(LogitsProcessor): def __init__(self, reps: int = 8, max_period: int = 4): self.reps, self.max_period = reps, max_period def __call__(self, input_ids, scores): for b in range(input_ids.shape[0]): tail = input_ids[b, -self.reps * self.max_period:].tolist() for p in range(1, self.max_period + 1): need = p * self.reps if len(tail) >= need and tail[-need:] == tail[-p:] * self.reps: scores[b, tail[-p]] = -float("inf") break return scores class TokenBias(LogitsProcessor): def __init__(self, ids: list[int], bias: float): self.ids, self.bias = ids, bias def __call__(self, input_ids, scores): scores[:, self.ids] += self.bias return scores class Transcriber: def __init__(self, model_dir: str, device: str = "cuda", num_beams: int = 1, punct_bias: float = 0.0, fallback: bool = True, max_cr: float = 1.8, gate_db: float | None = -55.0): self.dev = device self.dtype = torch.bfloat16 if device == "cuda" else torch.float32 self.model = AutoModelForCausalLM.from_pretrained(model_dir, trust_remote_code=True, dtype=self.dtype, attn_implementation="sdpa").to(device).eval() self.proc = AutoProcessor.from_pretrained(model_dir, trust_remote_code=True) self.num_beams, self.punct_bias, self.fallback, self.max_cr, self.gate_db = num_beams, punct_bias, fallback, max_cr, gate_db tok = self.proc.tokenizer self.punct_ids = sorted({ids[0] for s in ("！", "？", "!", "?", "っ", "ッ", "っ！", "！！", "！？", "ッ！") if len(ids := tok.encode(s, add_special_tokens=False)) == 1}) def _amp(self): return torch.autocast("cuda", dtype=torch.bfloat16) if self.dev == "cuda" else contextlib.nullcontext() def _generate(self, inputs, dur: float, temp: float) -> str: gc = copy.deepcopy(self.model.generation_config) gc.do_sample = temp > 0 if temp > 0: gc.temperature, gc.top_p, gc.top_k = temp, 1.0, 0 elif self.num_beams > 1: gc.num_beams = self.num_beams gc.max_new_tokens = min(4096, int(dur * 20) + 64) procs = [LoopGuard()] + ([TokenBias(self.punct_ids, self.punct_bias)] if self.punct_bias else []) with torch.inference_mode(), self._amp(): out = self.model.generate(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], input_features=inputs["input_features"], audio_feature_lengths=inputs["audio_feature_lengths"], audio_chunk_mapping=inputs["audio_chunk_mapping"], generation_config=gc, logits_processor=LogitsProcessorList(procs)) prompt_len = int(inputs["attention_mask"][0].sum().item()) return self.proc.tokenizer.decode(out[0][prompt_len:], skip_special_tokens=True).strip() def window(self, audio: np.ndarray, tmp_path: Path) -> list[tuple[float, float, str, str]]: """Transcribe one window (<= ~90 s works best). Returns (start, end, speaker, text).""" frames = frame_db(audio) if self.gate_db is not None and frames.max(initial=-120.0) < self.gate_db: return [] sf.write(tmp_path, audio, SR, subtype="PCM_16") with self._amp(): inputs = prepare_inputs(self.proc, build_transcription_messages(str(tmp_path)), max_length=131072, device=torch.device(self.dev)).to(self.dev) dur = len(audio) / SR best = None for temp in (0.0,) + ((0.2, 0.4, 0.6, 0.8) if self.fallback else ()): text = self._generate(inputs, dur, temp) c = compression_ratio(text) if best is None or c < best[0]: best = (c, text) if c <= self.max_cr: break segs = parse(best[1]) if self.gate_db is not None: hop = 0.05 segs = [s for s in segs if frames[int(s[0] / hop):max(int(s[1] / hop), int(s[0] / hop) + 1)].max(initial=-120.0) >= self.gate_db] return segs def __call__(self, path: Path, window_sec: float = 90.0) -> list[dict]: audio = load(path) tmp = Path(tempfile.gettempdir()) / f"moss_asr_{os.getpid()}.wav" # the processor reads audio from a path out = [] try: for off, chunk in split_long(audio, window_sec): for s, e, spk, txt in self.window(chunk, tmp): out.append({"start": round(off + s, 2), "end": round(off + e, 2), "speaker": spk, "text": txt, "window_start": round(off, 2)}) finally: tmp.unlink(missing_ok=True) return out def split_long(audio: np.ndarray, window_sec: float): """Cut at the quietest 50 ms frame in the last third of each window.""" n, win = len(audio), int(window_sec * SR) pos = 0 while pos < n: if n - pos <= win: yield pos / SR, audio[pos:] return seg = audio[pos:pos + win] db = frame_db(seg) lo = len(db) * 2 // 3 cut = (lo + int(np.argmin(db[lo:]))) * int(0.05 * SR) yield pos / SR, audio[pos:pos + cut] pos += cut def srt_time(t: float) -> str: ms = int(round(t * 1000)) return f"{ms // 3600000:02d}:{ms // 60000 % 60:02d}:{ms // 1000 % 60:02d},{ms % 1000:03d}" def main(): ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("inputs", nargs="+", help="audio files or folders") ap.add_argument("--model", default=str(HERE.parent)) ap.add_argument("--out", help="write JSON lines (one per file) here") ap.add_argument("--srt", help="write an SRT subtitle file (single input only)") ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") ap.add_argument("--num_beams", type=int, default=1) ap.add_argument("--punct_bias", type=float, default=0.0, help="logit bonus for ！？っ (sentence-final marks)") ap.add_argument("--no_fallback", action="store_true") ap.add_argument("--gate_db", type=float, default=-55.0) ap.add_argument("--no_gate", action="store_true") ap.add_argument("--window", type=float, default=90.0, help="max seconds per decoding window") args = ap.parse_args() files = [] for p in map(Path, args.inputs): files += sorted(f for f in p.rglob("*") if f.suffix.lower() in AUDIO_EXT) if p.is_dir() else [p] tr = Transcriber(args.model, args.device, args.num_beams, args.punct_bias, not args.no_fallback, gate_db=None if args.no_gate else args.gate_db) fout = open(args.out, "w", encoding="utf-8") if args.out else None for f in files: segs = tr(f, args.window) if fout: fout.write(json.dumps({"audio": str(f), "text": "".join(s["text"] for s in segs), "segments": segs}, ensure_ascii=False) + "\n") fout.flush() print(f"== {f}") for s in segs: print(f"[{s['start']:7.2f}-{s['end']:7.2f}] {s['speaker']}: {s['text']}") if args.srt: with open(args.srt, "w", encoding="utf-8") as fs: for i, s in enumerate(segs, 1): fs.write(f"{i}\n{srt_time(s['start'])} --> {srt_time(s['end'])}\n{s['speaker']}: {s['text']}\n\n") if fout: fout.close() if __name__ == "__main__": main()