Blame for code/transcribe.py

"""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()