Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
97 lines
4.5 KiB
Python
97 lines
4.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Run the pinned upstream implementation without changing its generation path."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import statistics
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import soundfile as sf
|
|
import torch
|
|
from transformers import AutoModelForCausalLM, AutoProcessor
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--model", type=Path, required=True)
|
|
parser.add_argument("--upstream", type=Path, required=True)
|
|
parser.add_argument("--audio", type=Path, nargs="+", required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--log", type=Path, required=True)
|
|
parser.add_argument("--dtype", choices=("bf16", "f32"), default="bf16")
|
|
parser.add_argument("--max-tokens", type=int, default=2048)
|
|
parser.add_argument("--repeats", type=int, default=3)
|
|
args = parser.parse_args()
|
|
if args.repeats < 1 or args.max_tokens < 1:
|
|
parser.error("repeats and max-tokens must be positive")
|
|
args.log.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
logging.basicConfig(level=logging.INFO, handlers=[
|
|
logging.FileHandler(args.log, mode="w"), logging.StreamHandler()])
|
|
sys.path.insert(0, str(args.upstream.resolve()))
|
|
from moss_transcribe_diarize import parse_transcript
|
|
from moss_transcribe_diarize.inference_utils import build_transcription_messages, generate_transcription
|
|
|
|
torch.set_num_threads(8)
|
|
torch.manual_seed(0)
|
|
torch.backends.cuda.matmul.allow_tf32 = False
|
|
torch.backends.cudnn.allow_tf32 = False
|
|
device = torch.device("cuda:0")
|
|
dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32
|
|
logging.info("torch=%s device=%s dtype=%s TF32=false", torch.__version__,
|
|
torch.cuda.get_device_name(device), args.dtype)
|
|
started = time.perf_counter()
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
str(args.model), trust_remote_code=True, local_files_only=True,
|
|
dtype=dtype, attn_implementation="sdpa").to(device).eval()
|
|
processor = AutoProcessor.from_pretrained(
|
|
str(args.model), trust_remote_code=True, local_files_only=True)
|
|
torch.cuda.synchronize()
|
|
logging.info("load_seconds=%.6f weight_bytes=%d resident_allocated_bytes=%d",
|
|
time.perf_counter() - started,
|
|
sum(p.numel() * p.element_size() for p in model.parameters()),
|
|
torch.cuda.memory_allocated())
|
|
|
|
results = []
|
|
for audio in args.audio:
|
|
info = sf.info(audio)
|
|
with audio.open("rb") as stream:
|
|
digest = hashlib.file_digest(stream, "sha256").hexdigest()
|
|
messages = build_transcription_messages(audio)
|
|
trials = []
|
|
for trial in range(args.repeats + 1):
|
|
torch.cuda.synchronize()
|
|
torch.cuda.reset_peak_memory_stats()
|
|
started = time.perf_counter()
|
|
result = generate_transcription(
|
|
model, processor, messages, max_new_tokens=args.max_tokens,
|
|
do_sample=False, device=device, dtype=dtype)
|
|
torch.cuda.synchronize()
|
|
elapsed = time.perf_counter() - started
|
|
if result["generated_tokens"] >= args.max_tokens:
|
|
raise RuntimeError(f"Generation hit max-tokens; increase it before using this as a baseline: {audio}")
|
|
result.update(seconds=elapsed, rtf=elapsed / info.duration,
|
|
peak_allocated_bytes=torch.cuda.max_memory_allocated(),
|
|
peak_reserved_bytes=torch.cuda.max_memory_reserved())
|
|
result["segments"] = [
|
|
dict(start=s.start, end=s.end, speaker=s.speaker, text=s.text)
|
|
for s in parse_transcript(result["text"])]
|
|
logging.info("audio=%s trial=%d result=%s", audio, trial, json.dumps(result, ensure_ascii=False))
|
|
trials.append(result)
|
|
if any(t["text"] != trials[0]["text"] for t in trials):
|
|
raise RuntimeError(f"Upstream greedy output changed across repeated runs: {audio}")
|
|
results.append(dict(audio=str(audio.resolve()), sha256=digest,
|
|
duration=info.duration, dtype=args.dtype,
|
|
max_tokens=args.max_tokens, trials=trials,
|
|
median_seconds=statistics.median(t["seconds"] for t in trials[1:])))
|
|
args.output.write_text(json.dumps(results, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|