Move model files to repo root

b85aa64 verified 15 days ago

10.3 kB

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

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