#!/usr/bin/env python3 """Audit an audiobook/EPUB pair before building a single-speaker TTS corpus.""" from __future__ import annotations import argparse import collections import difflib import json import re import subprocess import sys import xml.etree.ElementTree as ET import zipfile from pathlib import Path, PurePosixPath WORD_RE = re.compile(r"[а-яёa-z0-9]+", re.IGNORECASE) WHITESPACE_RE = re.compile(r"\s+") def normalized_words(text: str) -> list[str]: return [word.replace("ё", "е") for word in WORD_RE.findall(text.lower())] def element_text(root: ET.Element) -> str: ignored = {"script", "style", "svg", "math"} parts: list[str] = [] def visit(node: ET.Element) -> None: tag = node.tag.rsplit("}", 1)[-1].lower() if tag in ignored: return if node.text: parts.append(node.text) for child in node: visit(child) if child.tail: parts.append(child.tail) if tag in {"p", "div", "section", "h1", "h2", "h3", "h4", "h5", "h6", "li", "br"}: parts.append("\n") visit(root) lines = [WHITESPACE_RE.sub(" ", line).strip() for line in "".join(parts).splitlines()] return "\n".join(line for line in lines if line) def extract_epub(epub_path: Path, output_dir: Path) -> tuple[str, list[dict[str, object]]]: with zipfile.ZipFile(epub_path) as archive: container = ET.fromstring(archive.read("META-INF/container.xml")) rootfile = next( node.attrib["full-path"] for node in container.iter() if node.tag.rsplit("}", 1)[-1] == "rootfile" ) opf = ET.fromstring(archive.read(rootfile)) opf_dir = PurePosixPath(rootfile).parent manifest = { node.attrib["id"]: node.attrib["href"] for node in opf.iter() if node.tag.rsplit("}", 1)[-1] == "item" and "id" in node.attrib and "href" in node.attrib } spine_ids = [ node.attrib["idref"] for node in opf.iter() if node.tag.rsplit("}", 1)[-1] == "itemref" and "idref" in node.attrib ] sections: list[dict[str, object]] = [] texts: list[str] = [] for index, item_id in enumerate(spine_ids): href = manifest.get(item_id) if not href: continue member = str(opf_dir / PurePosixPath(href)) try: root = ET.fromstring(archive.read(member)) except (KeyError, ET.ParseError): continue text = element_text(root) if not text: continue headings = [ WHITESPACE_RE.sub(" ", "".join(node.itertext())).strip() for node in root.iter() if node.tag.rsplit("}", 1)[-1].lower() in {"h1", "h2", "h3"} ] sections.append( { "index": index, "href": member, "characters": len(text), "words": len(normalized_words(text)), "headings": [heading for heading in headings if heading], } ) texts.append(text) book_text = "\n\n".join(texts) output_dir.mkdir(parents=True, exist_ok=True) (output_dir / "book.txt").write_text(book_text, encoding="utf-8") (output_dir / "spine.json").write_text( json.dumps(sections, ensure_ascii=False, indent=2), encoding="utf-8" ) return book_text, sections def best_text_match(transcript: str, book_words: list[str]) -> dict[str, object]: spoken = normalized_words(transcript) if not spoken or not book_words: return {"ratio": 0.0, "book_excerpt": "", "word_offset": None} positions: dict[str, list[int]] = collections.defaultdict(list) for index, word in enumerate(book_words): if len(word) >= 4: positions[word].append(index) votes: collections.Counter[int] = collections.Counter() for spoken_index, word in enumerate(spoken): candidates = positions.get(word, ()) if len(candidates) <= 200: votes.update(book_index - spoken_index for book_index in candidates) offsets = [offset for offset, _ in votes.most_common(30)] or [0] best_ratio = 0.0 best_offset = 0 best_window: list[str] = [] for offset in offsets: start = max(0, offset - 8) window = book_words[start : start + len(spoken) + 16] ratio = difflib.SequenceMatcher(None, spoken, window, autojunk=False).ratio() if ratio > best_ratio: best_ratio, best_offset, best_window = ratio, start, window return { "ratio": round(best_ratio, 4), "book_excerpt": " ".join(best_window), "word_offset": best_offset, } def transcribe_samples( audio_files: list[Path], book_text: str, output_dir: Path, model_name: str, sample_seconds: int, ) -> list[dict[str, object]]: try: import av import imageio_ffmpeg from faster_whisper import WhisperModel except ImportError as error: raise SystemExit( "Install audit dependencies into PYTHONPATH: faster-whisper imageio-ffmpeg" ) from error chosen = [audio_files[0], audio_files[len(audio_files) // 2], audio_files[-1]] samples_dir = output_dir / "samples" samples_dir.mkdir(parents=True, exist_ok=True) ffmpeg = imageio_ffmpeg.get_ffmpeg_exe() model = WhisperModel( model_name, device="cpu", compute_type="int8", download_root=str(output_dir / "models"), ) book_words = normalized_words(book_text) results: list[dict[str, object]] = [] for source in chosen: with av.open(str(source)) as media: stream = media.streams.audio[0] duration = float(stream.duration * stream.time_base) if stream.duration is not None else 0.0 start = min(max(45.0, duration * 0.35), max(0.0, duration - sample_seconds - 5.0)) wav_path = samples_dir / f"{source.stem[:4]}-{int(start):05d}.wav" subprocess.run( [ ffmpeg, "-hide_banner", "-loglevel", "error", "-y", "-ss", f"{start:.3f}", "-i", str(source), "-t", str(sample_seconds), "-ac", "1", "-ar", "16000", str(wav_path), ], check=True, ) segments, info = model.transcribe( str(wav_path), language="ru", beam_size=5, vad_filter=True, condition_on_previous_text=True, ) transcript = " ".join(segment.text.strip() for segment in segments).strip() result = { "source": source.name, "source_duration_seconds": round(duration, 3), "sample_start_seconds": round(start, 3), "sample_duration_seconds": sample_seconds, "detected_language": info.language, "language_probability": round(info.language_probability, 4), "transcript": transcript, } result.update(best_text_match(transcript, book_words)) results.append(result) (output_dir / "alignment_samples.json").write_text( json.dumps(results, ensure_ascii=False, indent=2), encoding="utf-8" ) return results def main() -> int: parser = argparse.ArgumentParser() parser.add_argument("--source-dir", type=Path, required=True) parser.add_argument("--epub", type=Path) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--model", default="base") parser.add_argument("--sample-seconds", type=int, default=75) parser.add_argument("--skip-asr", action="store_true") args = parser.parse_args() audio_files = sorted(args.source_dir.glob("*.mp3")) epub_path = args.epub or next(args.source_dir.glob("*.epub"), None) if not audio_files: parser.error("No MP3 files found") if epub_path is None or not epub_path.is_file(): parser.error("EPUB file not found") book_text, sections = extract_epub(epub_path, args.output_dir) report: dict[str, object] = { "source_dir": str(args.source_dir.resolve()), "epub": str(epub_path.resolve()), "audio_files": [str(path.resolve()) for path in audio_files], "audio_file_count": len(audio_files), "epub_section_count": len(sections), "book_characters": len(book_text), "book_words": len(normalized_words(book_text)), } if not args.skip_asr: report["alignment_samples"] = transcribe_samples( audio_files, book_text, args.output_dir, args.model, args.sample_seconds ) (args.output_dir / "audit.json").write_text( json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8" ) print(json.dumps(report, ensure_ascii=False, indent=2)) return 0 if __name__ == "__main__": sys.exit(main())