365 lines
12 KiB
Python
365 lines
12 KiB
Python
#!/usr/bin/env python3
|
||
"""Align ASR sidecars to exact EPUB text and create a Piper WAV/metadata dataset."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import csv
|
||
import difflib
|
||
import json
|
||
import re
|
||
import subprocess
|
||
import sys
|
||
import wave
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
|
||
|
||
WORD_RE = re.compile(r"[а-яёa-z0-9]+", re.IGNORECASE)
|
||
|
||
|
||
def normalize_word(value: str) -> str:
|
||
return value.lower().replace("ё", "е")
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class BookWord:
|
||
value: str
|
||
start: int
|
||
end: int
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class AsrToken:
|
||
value: str
|
||
track: str
|
||
segment_id: int
|
||
start: float
|
||
end: float
|
||
|
||
|
||
@dataclass
|
||
class AlignedSegment:
|
||
track: str
|
||
source: Path
|
||
start: float
|
||
end: float
|
||
book_start: int
|
||
book_end: int
|
||
matched_tokens: int
|
||
total_tokens: int
|
||
avg_logprob: float
|
||
no_speech_prob: float
|
||
|
||
@property
|
||
def duration(self) -> float:
|
||
return self.end - self.start
|
||
|
||
@property
|
||
def coverage(self) -> float:
|
||
return self.matched_tokens / max(1, self.total_tokens)
|
||
|
||
|
||
def load_book(path: Path) -> tuple[str, list[BookWord]]:
|
||
text = path.read_text(encoding="utf-8")
|
||
words = [
|
||
BookWord(normalize_word(match.group()), match.start(), match.end())
|
||
for match in WORD_RE.finditer(text)
|
||
]
|
||
return text, words
|
||
|
||
|
||
def load_asr(sidecars: list[Path]) -> tuple[list[AsrToken], list[dict[str, object]], dict[str, Path]]:
|
||
tokens: list[AsrToken] = []
|
||
segments: list[dict[str, object]] = []
|
||
sources: dict[str, Path] = {}
|
||
for sidecar in sidecars:
|
||
payload = json.loads(sidecar.read_text(encoding="utf-8"))
|
||
track = sidecar.stem
|
||
source = Path(payload["source"])
|
||
sources[track] = source
|
||
for segment in payload["segments"]:
|
||
token_start = len(tokens)
|
||
for word in segment.get("words", []):
|
||
normalized = [normalize_word(match.group()) for match in WORD_RE.finditer(word["word"])]
|
||
for value in normalized:
|
||
tokens.append(
|
||
AsrToken(
|
||
value=value,
|
||
track=track,
|
||
segment_id=int(segment["id"]),
|
||
start=float(word["start"]),
|
||
end=float(word["end"]),
|
||
)
|
||
)
|
||
segments.append(
|
||
{
|
||
"track": track,
|
||
"source": source,
|
||
"segment_id": int(segment["id"]),
|
||
"start": float(segment["start"]),
|
||
"end": float(segment["end"]),
|
||
"avg_logprob": float(segment["avg_logprob"]),
|
||
"no_speech_prob": float(segment["no_speech_prob"]),
|
||
"token_start": token_start,
|
||
"token_end": len(tokens),
|
||
}
|
||
)
|
||
return tokens, segments, sources
|
||
|
||
|
||
def token_mapping(asr_tokens: list[AsrToken], book_words: list[BookWord]) -> dict[int, int]:
|
||
matcher = difflib.SequenceMatcher(
|
||
None,
|
||
[token.value for token in asr_tokens],
|
||
[word.value for word in book_words],
|
||
autojunk=True,
|
||
)
|
||
mapping: dict[int, int] = {}
|
||
for block in matcher.get_matching_blocks():
|
||
for offset in range(block.size):
|
||
mapping[block.a + offset] = block.b + offset
|
||
return mapping
|
||
|
||
|
||
def align_segments(
|
||
segments: list[dict[str, object]],
|
||
mapping: dict[int, int],
|
||
book_words: list[BookWord],
|
||
) -> list[AlignedSegment]:
|
||
aligned: list[AlignedSegment] = []
|
||
previous_book_end = -1
|
||
for segment in segments:
|
||
start_index = int(segment["token_start"])
|
||
end_index = int(segment["token_end"])
|
||
mapped = [mapping[index] for index in range(start_index, end_index) if index in mapping]
|
||
if not mapped:
|
||
continue
|
||
first, last = min(mapped), max(mapped)
|
||
if first < previous_book_end:
|
||
continue
|
||
asr_count = max(1, end_index - start_index)
|
||
book_count = last - first + 1
|
||
if book_count > asr_count * 1.8 + 8:
|
||
continue
|
||
item = AlignedSegment(
|
||
track=str(segment["track"]),
|
||
source=Path(segment["source"]),
|
||
start=float(segment["start"]),
|
||
end=float(segment["end"]),
|
||
book_start=book_words[first].start,
|
||
book_end=book_words[last].end,
|
||
matched_tokens=len(mapped),
|
||
total_tokens=asr_count,
|
||
avg_logprob=float(segment["avg_logprob"]),
|
||
no_speech_prob=float(segment["no_speech_prob"]),
|
||
)
|
||
aligned.append(item)
|
||
previous_book_end = last
|
||
return aligned
|
||
|
||
|
||
def merge_segments(items: list[AlignedSegment], book_text: str, max_duration: float) -> list[AlignedSegment]:
|
||
merged: list[AlignedSegment] = []
|
||
current: AlignedSegment | None = None
|
||
for item in items:
|
||
eligible = item.coverage >= 0.55 and item.avg_logprob >= -1.2 and item.no_speech_prob <= 0.5
|
||
if not eligible or item.duration <= 0:
|
||
if current is not None:
|
||
merged.append(current)
|
||
current = None
|
||
continue
|
||
if current is None:
|
||
current = item
|
||
continue
|
||
gap = item.start - current.end
|
||
combined_duration = item.end - current.start
|
||
same_track = item.track == current.track and item.source == current.source
|
||
near_in_book = 0 <= item.book_start - current.book_end <= 120
|
||
if same_track and gap <= 0.8 and near_in_book and combined_duration <= max_duration:
|
||
total = current.total_tokens + item.total_tokens
|
||
current.end = item.end
|
||
current.book_end = item.book_end
|
||
current.matched_tokens += item.matched_tokens
|
||
current.avg_logprob = (
|
||
current.avg_logprob * current.total_tokens + item.avg_logprob * item.total_tokens
|
||
) / total
|
||
current.no_speech_prob = max(current.no_speech_prob, item.no_speech_prob)
|
||
current.total_tokens = total
|
||
text = book_text[current.book_start : current.book_end].rstrip()
|
||
if current.duration >= 3.0 and text.endswith((".", "!", "?", "…", ":", ";")):
|
||
merged.append(current)
|
||
current = None
|
||
else:
|
||
merged.append(current)
|
||
current = item
|
||
if current is not None:
|
||
merged.append(current)
|
||
return merged
|
||
|
||
|
||
def clean_label(value: str) -> str:
|
||
return re.sub(r"\s+", " ", value).strip(" —–-\t\r\n")
|
||
|
||
|
||
def is_valid_wav(path: Path, sample_rate: int) -> bool:
|
||
if not path.is_file():
|
||
return False
|
||
try:
|
||
with wave.open(str(path), "rb") as audio:
|
||
return (
|
||
audio.getnchannels() == 1
|
||
and audio.getsampwidth() == 2
|
||
and audio.getframerate() == sample_rate
|
||
and audio.getnframes() > 0
|
||
)
|
||
except (EOFError, wave.Error):
|
||
return False
|
||
|
||
|
||
def create_dataset(
|
||
items: list[AlignedSegment],
|
||
book_text: str,
|
||
output_dir: Path,
|
||
ffmpeg: str,
|
||
sample_rate: int,
|
||
min_duration: float,
|
||
max_duration: float,
|
||
min_coverage: float,
|
||
) -> dict[str, object]:
|
||
wav_dir = output_dir / "wav"
|
||
wav_dir.mkdir(parents=True, exist_ok=True)
|
||
metadata_path = output_dir / "metadata.csv"
|
||
rows: list[tuple[str, str]] = []
|
||
accepted_seconds = 0.0
|
||
rejected = 0
|
||
details: list[dict[str, object]] = []
|
||
for index, item in enumerate(items):
|
||
text = clean_label(book_text[item.book_start : item.book_end])
|
||
duration = item.duration
|
||
if not (min_duration <= duration <= max_duration) or item.coverage < min_coverage:
|
||
rejected += 1
|
||
continue
|
||
if len(text) < 8 or len(text) > 320 or "|" in text:
|
||
rejected += 1
|
||
continue
|
||
clip_name = f"aletheia_ru_{len(rows):06d}.wav"
|
||
clip_path = wav_dir / clip_name
|
||
start = max(0.0, item.start - 0.06)
|
||
end = item.end + 0.08
|
||
if not is_valid_wav(clip_path, sample_rate):
|
||
subprocess.run(
|
||
[
|
||
ffmpeg,
|
||
"-hide_banner",
|
||
"-loglevel",
|
||
"error",
|
||
"-y",
|
||
"-ss",
|
||
f"{start:.3f}",
|
||
"-i",
|
||
str(item.source),
|
||
"-t",
|
||
f"{end - start:.3f}",
|
||
"-ac",
|
||
"1",
|
||
"-ar",
|
||
str(sample_rate),
|
||
"-sample_fmt",
|
||
"s16",
|
||
str(clip_path),
|
||
],
|
||
check=True,
|
||
)
|
||
rows.append((clip_name, text))
|
||
accepted_seconds += duration
|
||
details.append(
|
||
{
|
||
"clip": clip_name,
|
||
"source": str(item.source),
|
||
"start": round(item.start, 3),
|
||
"end": round(item.end, 3),
|
||
"duration": round(duration, 3),
|
||
"coverage": round(item.coverage, 4),
|
||
"text": text,
|
||
}
|
||
)
|
||
with metadata_path.open("w", encoding="utf-8", newline="") as output:
|
||
writer = csv.writer(output, delimiter="|", lineterminator="\n")
|
||
writer.writerows(rows)
|
||
report = {
|
||
"schema": 1,
|
||
"clips": len(rows),
|
||
"accepted_seconds": round(accepted_seconds, 3),
|
||
"accepted_hours": round(accepted_seconds / 3600, 4),
|
||
"rejected_candidates": rejected,
|
||
"sample_rate": sample_rate,
|
||
"min_duration": min_duration,
|
||
"max_duration": max_duration,
|
||
"min_coverage": min_coverage,
|
||
"items": details,
|
||
}
|
||
(output_dir / "dataset_report.json").write_text(
|
||
json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8"
|
||
)
|
||
return report
|
||
|
||
|
||
def main() -> int:
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--book-text", type=Path, required=True)
|
||
parser.add_argument("--transcripts-dir", type=Path, required=True)
|
||
parser.add_argument("--output-dir", type=Path, required=True)
|
||
parser.add_argument("--sample-rate", type=int, default=22050)
|
||
parser.add_argument("--min-duration", type=float, default=2.0)
|
||
parser.add_argument("--max-duration", type=float, default=12.0)
|
||
parser.add_argument("--min-coverage", type=float, default=0.62)
|
||
parser.add_argument("--dry-run", action="store_true")
|
||
args = parser.parse_args()
|
||
|
||
try:
|
||
import imageio_ffmpeg
|
||
except ImportError as error:
|
||
raise SystemExit("imageio-ffmpeg is missing from PYTHONPATH") from error
|
||
|
||
sidecars = sorted(args.transcripts_dir.glob("*.json"))
|
||
if not sidecars:
|
||
parser.error("No transcript sidecars found")
|
||
book_text, book_words = load_book(args.book_text)
|
||
asr_tokens, segments, _ = load_asr(sidecars)
|
||
mapping = token_mapping(asr_tokens, book_words)
|
||
aligned = align_segments(segments, mapping, book_words)
|
||
merged = merge_segments(aligned, book_text, args.max_duration)
|
||
summary = {
|
||
"transcript_files": len(sidecars),
|
||
"book_words": len(book_words),
|
||
"asr_tokens": len(asr_tokens),
|
||
"exact_token_matches": len(mapping),
|
||
"exact_token_match_ratio": round(len(mapping) / max(1, len(asr_tokens)), 4),
|
||
"aligned_segments": len(aligned),
|
||
"merged_candidates": len(merged),
|
||
}
|
||
if args.dry_run:
|
||
print(json.dumps(summary, ensure_ascii=False, indent=2))
|
||
return 0
|
||
report = create_dataset(
|
||
merged,
|
||
book_text,
|
||
args.output_dir,
|
||
imageio_ffmpeg.get_ffmpeg_exe(),
|
||
args.sample_rate,
|
||
args.min_duration,
|
||
args.max_duration,
|
||
args.min_coverage,
|
||
)
|
||
report.update(summary)
|
||
(args.output_dir / "dataset_report.json").write_text(
|
||
json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8"
|
||
)
|
||
print(json.dumps({key: value for key, value in report.items() if key != "items"}, ensure_ascii=False, indent=2))
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|