readalign 1.0.0__py3-none-any.whl

readalign/fs.py

@@ -0,0 +1,22 @@
+"""Filesystem writes that fail with a sentence instead of a traceback."""
+
+from __future__ import annotations
+
+import os
+
+from .errors import InputError
+
+
+def ensure_directory(path: str, what: str = "output directory") -> None:
+    try:
+        os.makedirs(path, exist_ok=True)
+    except OSError as exc:
+        raise InputError(f"cannot create the {what} {path}: {exc}") from exc
+
+
+def write_text(path: str, text: str) -> None:
+    try:
+        with open(path, "w", encoding="utf-8", newline="\n") as handle:
+            handle.write(text)
+    except OSError as exc:
+        raise InputError(f"cannot write {path}: {exc}") from exc

readalign/members.py

@@ -0,0 +1,72 @@
+"""What an output archive member can be, and how to read one.
+
+An audiobook is three orders of magnitude larger than everything else readalign handles, so
+members it does not need to look at are carried as a reference and streamed straight through to
+the output rather than held as bytes.
+"""
+
+from __future__ import annotations
+
+import os
+import zipfile
+from contextlib import AbstractContextManager
+from dataclasses import dataclass
+from typing import BinaryIO
+
+from .errors import InputError
+
+
+@dataclass(frozen=True)
+class SourceFile:
+    """A member whose bytes are a file on disk, such as a prepared audio track."""
+
+    path: str
+
+
+@dataclass(frozen=True)
+class ArchiveMember:
+    """A member copied through from the input EPUB without being parsed or changed."""
+
+    epub_path: str
+    name: str
+
+
+#: An output archive as member name to content.
+Member = dict[str, "bytes | SourceFile | ArchiveMember"]
+
+
+def open_archive(path: str) -> zipfile.ZipFile:
+    """Open an EPUB for reading, turning every unreadable-path case into a clear message."""
+    try:
+        return zipfile.ZipFile(path)
+    except FileNotFoundError:
+        raise InputError(f"no such file: {path}") from None
+    except (IsADirectoryError, PermissionError) as exc:
+        # Windows raises PermissionError for a directory, where POSIX raises IsADirectoryError.
+        if os.path.isdir(path):
+            raise InputError(f"not an .epub file but a directory: {path}") from None
+        raise InputError(f"cannot read {path}: {exc.strerror or exc}") from None
+    except zipfile.BadZipFile as exc:
+        raise InputError(f"{path} is not a readable EPUB (not a zip archive): {exc}") from None
+    except OSError as exc:
+        raise InputError(f"cannot read {path}: {exc.strerror or exc}") from None
+
+class MemberReader(AbstractContextManager):
+    """Opens by-reference members, holding each source archive open only once."""
+
+    def __init__(self) -> None:
+        self._archives: dict[str, zipfile.ZipFile] = {}
+
+    def open(self, member: SourceFile | ArchiveMember) -> BinaryIO:
+        if isinstance(member, SourceFile):
+            return open(member.path, "rb")
+        archive = self._archives.get(member.epub_path)
+        if archive is None:
+            archive = zipfile.ZipFile(member.epub_path)
+            self._archives[member.epub_path] = archive
+        return archive.open(member.name)
+
+    def __exit__(self, *_) -> None:
+        while self._archives:
+            _, archive = self._archives.popitem()
+            archive.close()

readalign/package.py

@@ -0,0 +1,287 @@
+"""Rewrite the package document and write the finished EPUB.
+
+The OPF gains the audio and overlay manifest items, the ``media-overlay`` attribute on every
+narrated content document, and the ``media:`` metadata EPUB 3.3 section 9.3.5 requires. The
+archive is written mimetype-first and uncompressed, as OCF demands.
+"""
+
+from __future__ import annotations
+
+import os
+import posixpath
+import shutil
+import zipfile
+from urllib.parse import unquote
+
+from lxml import etree
+
+from .align import TimedSentence
+from .audio import AudioTrack
+from .clock import format_clock
+from .epub import OPF_NS, EpubPackage, local_name
+from .errors import InputError
+from .fs import ensure_directory
+from .members import Member, MemberReader, SourceFile
+from .smil import build_smil, overlay_duration
+
+AUDIO_DIR = "audio"
+OVERLAY_DIR = "media-overlays"
+
+ACTIVE_CLASS = "-epub-media-overlay-active"
+PLAYBACK_ACTIVE_CLASS = "-epub-media-overlay-playback-active"
+
+#: Stylesheet injected so the highlight is visible in readers that honour media:active-class.
+ACTIVE_CLASS_CSS = f""".{ACTIVE_CLASS} {{
+\tbackground-color: rgba(255, 224, 96, 0.55);
+\tcolor: inherit;
+}}
+"""
+
+_MEDIA_TYPES_BY_SUFFIX = {
+    ".mp3": "audio/mpeg",
+    ".m4a": "audio/mp4",
+    ".opus": "audio/ogg; codecs=opus",
+}
+
+
+def _manifest_element(package: EpubPackage) -> etree._Element:
+    root = package.opf_tree.getroot()
+    manifest = root.find(f"{{{OPF_NS}}}manifest")
+    if manifest is None:
+        raise InputError("the package document has no <manifest>")
+    return manifest
+
+
+def _metadata_element(package: EpubPackage) -> etree._Element:
+    root = package.opf_tree.getroot()
+    metadata = root.find(f"{{{OPF_NS}}}metadata")
+    if metadata is None:
+        raise InputError("the package document has no <metadata>")
+    return metadata
+
+
+def _free_id(package: EpubPackage, base: str) -> str:
+    if base not in package.manifest:
+        return base
+    suffix = 2
+    while f"{base}-{suffix}" in package.manifest:
+        suffix += 1
+    return f"{base}-{suffix}"
+
+
+def _add_manifest_item(
+    package: EpubPackage, manifest: etree._Element, item_id: str, href: str, media_type: str
+) -> etree._Element:
+    element = etree.SubElement(manifest, f"{{{OPF_NS}}}item")
+    element.set("id", item_id)
+    element.set("href", href)
+    element.set("media-type", media_type)
+    return element
+
+
+def _strip_existing_overlays(package: EpubPackage, manifest: etree._Element) -> set[str]:
+    """Drop overlays from an EPUB that already has them, so a re-run replaces rather than stacks."""
+    removed: set[str] = set()
+    for item in list(package.manifest.values()):
+        if item.media_type == "application/smil+xml":
+            manifest.remove(item.element)
+            package.manifest.pop(item.id, None)
+            removed.add(package.resolve(item.href))
+        elif item.element.get("media-overlay"):
+            del item.element.attrib["media-overlay"]
+    return removed
+
+
+def _drop_existing_overlay_metadata(metadata: etree._Element) -> None:
+    for element in list(metadata):
+        if element.tag == f"{{{OPF_NS}}}meta" and (element.get("property") or "").startswith(
+            "media:"
+        ):
+            metadata.remove(element)
+
+
+def _set_meta(metadata: etree._Element, prop: str, value: str, refines: str | None = None) -> None:
+    element = etree.SubElement(metadata, f"{{{OPF_NS}}}meta")
+    element.set("property", prop)
+    if refines:
+        element.set("refines", refines)
+    element.text = value
+
+
+def _links_to(doc, css_path: str) -> list[etree._Element]:
+    """The <link> elements in one document that point at ``css_path``."""
+    found = []
+    for element in doc.tree.getroot().iter():
+        if local_name(element.tag) != "link":
+            continue
+        href = element.get("href")
+        if not href:
+            continue
+        resolved = posixpath.normpath(
+            posixpath.join(posixpath.dirname(doc.zip_path), unquote(href))
+        )
+        if resolved == css_path:
+            found.append(element)
+    return found
+
+
+def _strip_previous_stylesheet(
+    package: EpubPackage, manifest: etree._Element, css_path: str
+) -> None:
+    """Drop an earlier run's stylesheet, so re-running replaces the link rather than stacking."""
+    for item in list(package.manifest.values()):
+        if item.media_type == "text/css" and package.resolve(item.href) == css_path:
+            manifest.remove(item.element)
+            package.manifest.pop(item.id, None)
+    for doc in package.docs:
+        for link in _links_to(doc, css_path):
+            link.getparent().remove(link)
+
+
+def _inject_stylesheet(
+    package: EpubPackage, manifest: etree._Element, href: str, css_path: str
+) -> None:
+    _strip_previous_stylesheet(package, manifest, css_path)
+    _add_manifest_item(package, manifest, _free_id(package, "readalign-css"), href, "text/css")
+    for doc in package.docs:
+        root = doc.tree.getroot()
+        head = next(
+            (element for element in root.iter() if local_name(element.tag) == "head"),
+            None,
+        )
+        if head is None:
+            continue
+        namespace = root.tag[: root.tag.index("}") + 1] if root.tag.startswith("{") else ""
+        link = etree.SubElement(head, f"{namespace}link")
+        link.set("rel", "stylesheet")
+        link.set("type", "text/css")
+        link.set("href", posixpath.relpath(css_path, posixpath.dirname(doc.zip_path)))
+
+
+def build_output(
+    package: EpubPackage,
+    entries_by_doc: dict[str, list[TimedSentence]],
+    tracks: list[AudioTrack],
+    *,
+    narrator: str | None = None,
+) -> tuple[Member, dict]:
+    """Produce the member-name-to-bytes map of the output EPUB, and a summary of what changed."""
+    manifest = _manifest_element(package)
+    metadata = _metadata_element(package)
+    _drop_existing_overlay_metadata(metadata)
+    replaced = _strip_existing_overlays(package, manifest)
+
+    by_index = {track.index: track for track in tracks}
+    used_tracks = sorted(
+        {entry.track_index for entries in entries_by_doc.values() for entry in entries}
+    )
+    audio_paths: dict[int, str] = {}
+    audio_hrefs: dict[int, str] = {}
+    files: Member = {
+        name: data for name, data in package.files.items() if name not in replaced
+    }
+
+    for index in used_tracks:
+        track = by_index[index]
+        suffix = os.path.splitext(track.path)[1].lower()
+        href = posixpath.join(AUDIO_DIR, f"part{index:03d}{suffix}")
+        zip_path = package.resolve(href)
+        audio_paths[index] = zip_path
+        audio_hrefs[index] = href
+        files[zip_path] = SourceFile(track.path)
+        _add_manifest_item(
+            package,
+            manifest,
+            _free_id(package, f"readalign-audio-{index:03d}"),
+            href,
+            _MEDIA_TYPES_BY_SUFFIX.get(suffix, track.media_type),
+        )
+
+    css_href = posixpath.join("css", "readalign.css")
+    css_path = package.resolve(css_href)
+    files[css_path] = ACTIVE_CLASS_CSS.encode("utf-8")
+    _inject_stylesheet(package, manifest, css_href, css_path)
+
+    overlays: list[tuple[str, float]] = []
+    total = 0.0
+    for doc in package.docs:
+        entries = entries_by_doc.get(doc.manifest_id) or []
+        if not entries:
+            continue
+        stem = posixpath.splitext(posixpath.basename(doc.href))[0]
+        href = posixpath.join(OVERLAY_DIR, f"{stem}.smil")
+        smil_path = package.resolve(href)
+        files[smil_path] = build_smil(smil_path, doc.zip_path, entries, by_index, audio_paths)
+        overlay_id = _free_id(package, f"readalign-smil-{doc.manifest_id}")
+        _add_manifest_item(package, manifest, overlay_id, href, "application/smil+xml")
+        package.manifest[doc.manifest_id].element.set("media-overlay", overlay_id)
+        duration = overlay_duration(entries)
+        overlays.append((overlay_id, duration))
+        total += duration
+
+    for overlay_id, duration in overlays:
+        _set_meta(metadata, "media:duration", format_clock(duration), refines=f"#{overlay_id}")
+    _set_meta(metadata, "media:duration", format_clock(total))
+    _set_meta(metadata, "media:active-class", ACTIVE_CLASS)
+    _set_meta(metadata, "media:playback-active-class", PLAYBACK_ACTIVE_CLASS)
+    if narrator:
+        _set_meta(metadata, "media:narrator", narrator)
+
+    for doc in package.docs:
+        files[doc.zip_path] = doc.serialize()
+    files[package.opf_path] = etree.tostring(
+        package.opf_tree, xml_declaration=True, encoding="utf-8"
+    )
+
+    summary = {
+        "overlays": len(overlays),
+        "total_duration": format_clock(total),
+        "audio_files": len(audio_paths),
+        "audio_hrefs": audio_hrefs,
+        "active_class": ACTIVE_CLASS,
+    }
+    return files, summary
+
+
+def write_epub(path: str, files: Member, original_order: list[str]) -> None:
+    """Write an OCF archive: mimetype first and stored, everything else deflated."""
+    ordered: list[str] = []
+    seen = set()
+    for name in ["mimetype", "META-INF/container.xml", *original_order]:
+        if name in files and name not in seen:
+            ordered.append(name)
+            seen.add(name)
+    ordered.extend(sorted(name for name in files if name not in seen))
+
+    ensure_directory(os.path.dirname(os.path.abspath(path)))
+
+    try:
+        with zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED) as archive, MemberReader() as read:
+            for name in ordered:
+                data = files[name]
+                if name == "mimetype":
+                    # OCF requires the first entry to be an uncompressed "mimetype".
+                    info = zipfile.ZipInfo("mimetype")
+                    info.compress_type = zipfile.ZIP_STORED
+                    info.external_attr = 0o644 << 16
+                    archive.writestr(info, data)
+                    continue
+                info = zipfile.ZipInfo(name)
+                info.external_attr = 0o644 << 16
+                info.compress_type = (
+                    zipfile.ZIP_STORED
+                    if posixpath.splitext(name)[1].lower() in _MEDIA_TYPES_BY_SUFFIX
+                    else zipfile.ZIP_DEFLATED
+                )
+                if isinstance(data, bytes):
+                    archive.writestr(info, data)
+                else:
+                    with read.open(data) as source, archive.open(info, "w") as target:
+                        shutil.copyfileobj(source, target, 1 << 20)
+    except OSError as exc:
+        raise InputError(f"cannot write {path}: {exc}") from exc
+
+
+def ensure_mimetype(files: Member) -> None:
+    """Set the OCF mimetype. It is fixed for EPUB, so the input's copy is never carried over."""
+    files["mimetype"] = b"application/epub+zip"

readalign/refine.py

@@ -0,0 +1,175 @@
+"""Optional CTC forced-alignment pass that sharpens sentence boundaries.
+
+Whisper's word timestamps come from attention weights and are typically good to a few tenths of
+a second. A frame-level CTC aligner does better, so ``--refine`` runs torchaudio's MMS_FA bundle
+over each already-aligned stretch and replaces the boundaries it is confident about.
+
+The MMS-300m aligner is licensed cc-by-nc-4.0, which is not a licence that belongs in a default
+install path, so this module is never imported unless ``--refine`` is passed and the model is
+never downloaded otherwise.
+"""
+
+from __future__ import annotations
+
+from .align import AlignmentResult, TimedSentence
+from .audio import SAMPLE_RATE, AudioTrack, read_wav_segment
+from .errors import DependencyError
+from .textnorm import tokenize_words
+
+#: Seconds of audio handed to the aligner at once. The CTC lattice is time x tokens, so this
+#: bounds peak memory independently of how long the book is.
+CHUNK_SECONDS = 40.0
+
+#: Context added either side of a chunk so a boundary word is never cut in half.
+PAD_SECONDS = 0.5
+
+#: Boundaries are only replaced when the new value is within this many seconds of the old one.
+MAX_SHIFT = 2.0
+
+LICENCE_NOTE = (
+    "the MMS-300m forced aligner is licensed CC BY-NC 4.0 (non-commercial). "
+    "Output timings refined with it inherit that restriction."
+)
+
+
+def _load_bundle(device: str):
+    try:
+        import torch
+        import torchaudio
+    except ImportError as exc:
+        raise DependencyError(
+            "--refine needs the optional extra: pip install 'readalign[refine]'"
+        ) from exc
+    try:
+        bundle = torchaudio.pipelines.MMS_FA
+        model = bundle.get_model()
+    except Exception as exc:
+        message = str(exc)
+        raise DependencyError(
+            f"could not load the MMS forced aligner. It is downloaded once from "
+            f"download.pytorch.org, so this needs network access. Underlying error: {message}"
+        ) from exc
+    resolved = device
+    if device == "auto":
+        resolved = "cuda" if torch.cuda.is_available() else "cpu"
+    model = model.to(resolved).eval()
+    return torch, bundle, model, resolved
+
+
+def _chunks(entries: list[TimedSentence]) -> list[list[TimedSentence]]:
+    chunks: list[list[TimedSentence]] = []
+    current: list[TimedSentence] = []
+    for entry in entries:
+        if current and entry.end - current[0].start > CHUNK_SECONDS:
+            chunks.append(current)
+            current = []
+        current.append(entry)
+    if current:
+        chunks.append(current)
+    return chunks
+
+
+def _transcript(
+    entries: list[TimedSentence], vocabulary: set[str]
+) -> tuple[list[str], list[tuple[int, int]]]:
+    """Flatten the chunk's sentences into aligner words, keeping each sentence's word range."""
+    words: list[str] = []
+    spans: list[tuple[int, int]] = []
+    for entry in entries:
+        tokens, _ = tokenize_words(entry.sentence.words)
+        start = len(words)
+        for token in tokens:
+            cleaned = "".join(character for character in token if character in vocabulary)
+            if cleaned:
+                words.append(cleaned)
+        spans.append((start, len(words)))
+    return words, spans
+
+
+def refine(
+    result: AlignmentResult,
+    tracks: list[AudioTrack],
+    *,
+    device: str = "auto",
+    log=print,
+) -> dict:
+    """Replace Whisper's sentence boundaries with CTC ones wherever both agree roughly."""
+    torch, bundle, model, resolved = _load_bundle(device)
+    tokenizer = bundle.get_tokenizer()
+    aligner = bundle.get_aligner()
+    vocabulary = set(bundle.get_dict())
+    by_index = {track.index: track for track in tracks}
+    log(f"  refining boundaries with MMS_FA on {resolved}")
+
+    adjusted = 0
+    failed = 0
+    total_shift = 0.0
+    for track_index, track in sorted(by_index.items()):
+        if track.wav_path is None:
+            continue
+        entries = sorted(
+            (entry for entry in result.timed if entry.track_index == track_index),
+            key=lambda entry: entry.start,
+        )
+        if not entries:
+            continue
+        for chunk in _chunks(entries):
+            words, spans = _transcript(chunk, vocabulary)
+            if not words:
+                continue
+            begin = max(0.0, chunk[0].start - track.offset - PAD_SECONDS)
+            finish = min(track.duration, chunk[-1].end - track.offset + PAD_SECONDS)
+            if finish - begin < 0.05:
+                continue
+            samples = read_wav_segment(track.wav_path, begin, finish)
+            if samples.size < SAMPLE_RATE // 10:
+                continue
+            waveform = torch.from_numpy(samples).unsqueeze(0).to(resolved)
+            try:
+                with torch.inference_mode():
+                    emission, _ = model(waveform)
+                    token_spans = aligner(emission[0], tokenizer(words))
+            except Exception as exc:
+                failed += 1
+                if failed == 1:
+                    log(f"  warning: the refiner failed on one chunk and was skipped: {exc}")
+                continue
+            ratio = waveform.size(1) / emission.size(1)
+            origin = begin + track.offset
+            starts = [origin + span[0].start * ratio / SAMPLE_RATE for span in token_spans]
+            ends = [origin + span[-1].end * ratio / SAMPLE_RATE for span in token_spans]
+            for entry, (first, last) in zip(chunk, spans, strict=True):
+                if last <= first or last > len(starts):
+                    continue
+                new_start, new_end = starts[first], ends[last - 1]
+                if new_end <= new_start:
+                    continue
+                if abs(new_start - entry.start) > MAX_SHIFT or abs(new_end - entry.end) > MAX_SHIFT:
+                    continue
+                total_shift += abs(new_start - entry.start)
+                entry.start, entry.end = new_start, new_end
+                entry.interpolated = False
+                adjusted += 1
+
+    _repair_order(result.timed, by_index)
+    skipped = f", {failed} chunk(s) skipped" if failed else ""
+    log(f"  refined {adjusted} sentence boundaries{skipped}")
+    return {
+        "device": resolved,
+        "adjusted": adjusted,
+        "failed_chunks": failed,
+        "mean_shift": round(total_shift / adjusted, 3) if adjusted else 0.0,
+        "licence": LICENCE_NOTE,
+    }
+
+
+def _repair_order(timed: list[TimedSentence], tracks: dict[int, AudioTrack]) -> None:
+    """Keep the sequence monotonic and inside each track after boundaries moved."""
+    previous: float | None = None
+    for entry in timed:
+        track = tracks[entry.track_index]
+        entry.start = min(max(entry.start, track.offset), track.end)
+        entry.end = min(max(entry.end, entry.start + 0.05), track.end)
+        if previous is not None and entry.start < previous:
+            entry.start = min(previous, entry.end - 0.05)
+        previous = entry.end