117 lines
3.9 KiB
Python
117 lines
3.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Create resumable word-timestamp ASR sidecars for audiobook tracks."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
|
|
def sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as source:
|
|
for block in iter(lambda: source.read(1024 * 1024), b""):
|
|
digest.update(block)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def serialize_segment(segment) -> dict[str, object]:
|
|
return {
|
|
"id": segment.id,
|
|
"start": round(segment.start, 3),
|
|
"end": round(segment.end, 3),
|
|
"text": segment.text.strip(),
|
|
"avg_logprob": round(segment.avg_logprob, 5),
|
|
"compression_ratio": round(segment.compression_ratio, 5),
|
|
"no_speech_prob": round(segment.no_speech_prob, 5),
|
|
"words": [
|
|
{
|
|
"start": round(word.start, 3),
|
|
"end": round(word.end, 3),
|
|
"word": word.word,
|
|
"probability": round(word.probability, 5),
|
|
}
|
|
for word in (segment.words or ())
|
|
],
|
|
}
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--source-dir", type=Path, required=True)
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
|
parser.add_argument("--model", default="base")
|
|
parser.add_argument("--model-cache", type=Path, required=True)
|
|
parser.add_argument("--max-files", type=int)
|
|
parser.add_argument("--force", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
from faster_whisper import WhisperModel
|
|
except ImportError as error:
|
|
raise SystemExit("faster-whisper is missing from PYTHONPATH") from error
|
|
|
|
files = sorted(args.source_dir.glob("*.mp3"))
|
|
if args.max_files is not None:
|
|
files = files[: max(0, args.max_files)]
|
|
if not files:
|
|
parser.error("No MP3 files found")
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
model = WhisperModel(
|
|
args.model,
|
|
device="cpu",
|
|
compute_type="int8",
|
|
download_root=str(args.model_cache),
|
|
local_files_only=True,
|
|
)
|
|
|
|
completed = 0
|
|
for index, audio_path in enumerate(files, 1):
|
|
output_path = args.output_dir / f"{audio_path.stem[:4]}.json"
|
|
if output_path.is_file() and not args.force:
|
|
print(f"[{index}/{len(files)}] skip {audio_path.name}", flush=True)
|
|
completed += 1
|
|
continue
|
|
started = time.monotonic()
|
|
print(f"[{index}/{len(files)}] transcribe {audio_path.name}", flush=True)
|
|
segments_iter, info = model.transcribe(
|
|
str(audio_path),
|
|
language="ru",
|
|
beam_size=5,
|
|
vad_filter=True,
|
|
word_timestamps=True,
|
|
condition_on_previous_text=True,
|
|
)
|
|
segments = [serialize_segment(segment) for segment in segments_iter]
|
|
payload = {
|
|
"schema": 1,
|
|
"source": str(audio_path.resolve()),
|
|
"source_bytes": audio_path.stat().st_size,
|
|
"source_sha256": sha256(audio_path),
|
|
"language": info.language,
|
|
"language_probability": round(info.language_probability, 5),
|
|
"duration": round(info.duration, 3),
|
|
"duration_after_vad": round(info.duration_after_vad, 3),
|
|
"elapsed_seconds": round(time.monotonic() - started, 3),
|
|
"segments": segments,
|
|
}
|
|
temporary = output_path.with_suffix(".json.tmp")
|
|
temporary.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
temporary.replace(output_path)
|
|
completed += 1
|
|
print(
|
|
f"[{index}/{len(files)}] wrote {output_path.name}: "
|
|
f"{len(segments)} segments in {payload['elapsed_seconds']}s",
|
|
flush=True,
|
|
)
|
|
print(json.dumps({"files": len(files), "completed": completed}, ensure_ascii=False))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|