From 15c1220ba3588484dd87817d3e5c7949ef6b3624 Mon Sep 17 00:00:00 2001 From: JB Date: Fri, 10 Oct 2025 11:17:14 -0700 Subject: [PATCH] feat: Enhance chunking logic to include display text and preserve whitespace in sentences --- abogen/chunking.py | 99 ++++++++++++++++++++------ abogen/epub3/exporter.py | 42 +++++++++-- abogen/web/templates/reader_embed.html | 10 ++- tests/test_chunk_helpers.py | 20 ++++++ 4 files changed, 144 insertions(+), 27 deletions(-) diff --git a/abogen/chunking.py b/abogen/chunking.py index 40f500f..265a1e0 100644 --- a/abogen/chunking.py +++ b/abogen/chunking.py @@ -1,7 +1,8 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Dict, Iterable, Iterator, List, Literal, Optional +from typing import Dict, Iterable, Iterator, List, Literal, Optional, Tuple +from typing import Pattern import re @@ -31,6 +32,7 @@ class Chunk: voice: Optional[str] = None voice_profile: Optional[str] = None voice_formula: Optional[str] = None + display_text: Optional[str] = None def as_dict(self) -> Dict[str, object]: return { @@ -43,6 +45,7 @@ class Chunk: "voice": self.voice, "voice_profile": self.voice_profile, "voice_formula": self.voice_formula, + "display_text": self.display_text, } @@ -53,19 +56,21 @@ def _iter_paragraphs(text: str) -> Iterator[str]: yield normalized -def _iter_sentences(paragraph: str) -> Iterator[str]: +def _iter_sentences(paragraph: str) -> Iterator[Tuple[str, str]]: if not paragraph: return start = 0 for match in _SENTENCE_SPLIT_REGEX.finditer(paragraph): end = match.end() - candidate = paragraph[start:end].strip() + raw_segment = paragraph[start:end] + candidate = raw_segment.strip() if candidate: - yield candidate + yield candidate, raw_segment start = match.end() - tail = paragraph[start:].strip() + tail_raw = paragraph[start:] + tail = tail_raw.strip() if tail: - yield tail + yield tail, tail_raw def _normalize_whitespace(value: str) -> str: @@ -77,28 +82,32 @@ def _normalize_chunk_text(value: str) -> str: return _normalize_whitespace(normalized) -def _split_sentences(paragraph: str) -> List[str]: +def _split_sentences(paragraph: str) -> List[Tuple[str, str]]: sentences = list(_iter_sentences(paragraph)) if not sentences: return [] - merged: List[str] = [] - buffer: List[str] = [] + merged: List[Tuple[str, str]] = [] + buffer_norm: List[str] = [] + buffer_raw: List[str] = [] - for sentence in sentences: - if buffer: - buffer.append(sentence) + for normalized_sentence, raw_sentence in sentences: + if buffer_norm: + buffer_norm.append(normalized_sentence) + buffer_raw.append(raw_sentence) else: - buffer = [sentence] + buffer_norm = [normalized_sentence] + buffer_raw = [raw_sentence] - if _ABBREVIATION_END_RE.search(sentence.rstrip()): + if _ABBREVIATION_END_RE.search(normalized_sentence.rstrip()): continue - merged.append(" ".join(buffer)) - buffer = [] + merged.append((" ".join(buffer_norm), "".join(buffer_raw))) + buffer_norm = [] + buffer_raw = [] - if buffer: - merged.append(" ".join(buffer)) + if buffer_norm: + merged.append((" ".join(buffer_norm), "".join(buffer_raw))) return merged @@ -140,6 +149,7 @@ def chunk_text( ).as_dict() payload["normalized_text"] = _normalize_chunk_text(paragraph) chunks.append(payload) + _attach_display_text(text, chunks) return chunks # Sentence level – flatten paragraphs into individual sentences @@ -148,9 +158,9 @@ def chunk_text( normalized_para = _normalize_whitespace(paragraph) if not normalized_para: continue - sentences = _split_sentences(normalized_para) or [normalized_para] - for sent_local_index, sentence in enumerate(sentences): - normalized_sentence = _normalize_whitespace(sentence) + sentence_pairs = _split_sentences(paragraph) or [(normalized_para, paragraph)] + for sent_local_index, (normalized_sentence, raw_sentence) in enumerate(sentence_pairs): + normalized_sentence = _normalize_whitespace(normalized_sentence) if not normalized_sentence: continue chunk_id = f"{prefix}_p{para_index:04d}_s{sent_local_index:04d}" @@ -165,13 +175,58 @@ def chunk_text( voice_profile=voice_profile, voice_formula=voice_formula, ).as_dict() - payload["normalized_text"] = _normalize_chunk_text(sentence) + payload["normalized_text"] = _normalize_chunk_text(raw_sentence) + payload["display_text"] = raw_sentence chunks.append(payload) sentence_index += 1 + _attach_display_text(text, chunks) return chunks +_DISPLAY_PATTERN_CACHE: Dict[str, Pattern[str]] = {} + + +def _build_display_pattern(text: str) -> Pattern[str]: + cached = _DISPLAY_PATTERN_CACHE.get(text) + if cached is not None: + return cached + escaped = re.escape(text) + escaped = escaped.replace(r"\ ", r"\s+") + pattern = re.compile(r"(\s*" + escaped + r"\s*)", re.DOTALL) + _DISPLAY_PATTERN_CACHE[text] = pattern + return pattern + + +def _search_source_span(source: str, normalized: str, start: int) -> Optional[Tuple[int, int]]: + if not normalized: + return None + pattern = _build_display_pattern(normalized) + match = pattern.search(source, start) + if not match: + return None + return match.start(1), match.end(1) + + +def _attach_display_text(source: str, chunks: List[Dict[str, object]]) -> None: + if not source or not chunks: + return + cursor = 0 + for chunk in chunks: + candidate = str(chunk.get("display_text") or chunk.get("text") or "") + if not candidate: + continue + match = _search_source_span(source, candidate, cursor) + if match is None and cursor: + match = _search_source_span(source, candidate, 0) + if match is None: + chunk.setdefault("display_text", candidate) + continue + start, end = match + chunk["display_text"] = source[start:end] + cursor = end + + def build_chunks_for_chapters( chapters: Iterable[Dict[str, object]], *, diff --git a/abogen/epub3/exporter.py b/abogen/epub3/exporter.py index d9f371d..b032dfa 100644 --- a/abogen/epub3/exporter.py +++ b/abogen/epub3/exporter.py @@ -22,6 +22,7 @@ class ChunkOverlay: end: Optional[float] speaker_id: str voice: Optional[str] + level: Optional[str] = None @dataclass(slots=True) @@ -233,14 +234,19 @@ class EPUB3PackageBuilder: if chunk_entry is None and chapter_chunks and position < len(chapter_chunks): chunk_entry = chapter_chunks[position] + level = None if chunk_entry is None: text = self.extraction.chapters[chapter_index].text speaker_id = str(marker.get("speaker_id") or "narrator") voice = marker.get("voice") else: - text = str(chunk_entry.get("text") or "") + display_text = chunk_entry.get("display_text") + text = str(display_text or chunk_entry.get("text") or "") speaker_id = str(chunk_entry.get("speaker_id") or marker.get("speaker_id") or "narrator") voice = chunk_entry.get("voice") or chunk_entry.get("resolved_voice") or marker.get("voice") + level = chunk_entry.get("level") or None + if chunk_entry is None: + level = None normalized_id = _normalize_chunk_id(chunk_id) if chunk_id else None if not normalized_id: @@ -257,6 +263,7 @@ class EPUB3PackageBuilder: end=_safe_float(marker.get("end")), speaker_id=speaker_id, voice=str(voice) if voice else None, + level=str(level) if level else None, ) ) @@ -275,6 +282,16 @@ class EPUB3PackageBuilder: chunk_html = "\n".join(_render_chunk_html(chunk) for chunk in chapter.chunks) if not chunk_html: chunk_html = "

" + original_block = "" + if chapter.chunks: + original_text = "".join(chunk.text or "" for chunk in chapter.chunks) + if original_text: + safe_original = html.escape(original_text) + original_block = ( + " " + ) return ( "\n" @@ -288,10 +305,17 @@ class EPUB3PackageBuilder: "
\n" "

{title}

\n" " {chunks}\n" + "{original_block}" "
\n" " \n" "\n" - ).format(lang=language, title=title, index=chapter.index + 1, chunks=chunk_html) + ).format( + lang=language, + title=title, + index=chapter.index + 1, + chunks=chunk_html, + original_block=("" if not original_block else f"{original_block}\n"), + ) def _render_chapter_smil(self, chapter: ChapterDocument, audio_href: str) -> str: par_lines = [] @@ -625,13 +649,14 @@ def _render_chunk_html(chunk: ChunkOverlay) -> str: escaped_id = html.escape(chunk.id) speaker_attr = f" data-speaker=\"{html.escape(chunk.speaker_id)}\"" if chunk.speaker_id else "" voice_attr = f" data-voice=\"{html.escape(chunk.voice)}\"" if chunk.voice else "" + level_attr = f" data-level=\"{html.escape(chunk.level)}\"" if chunk.level else "" raw_text = chunk.text or "" escaped_text = html.escape(raw_text) if not escaped_text: escaped_text = " " body = escaped_text.replace("\n", "\n ") return ( - f"
\n" + f"
\n" f" {body}\n" "
" ) @@ -803,7 +828,16 @@ h1 { } div.chunk { - margin-bottom: 1em; + margin: 0 0 1em 0; + white-space: pre-wrap; +} + +div.chunk[data-level="sentence"] { + margin-bottom: 0.5em; +} + +div.chunk[data-level="sentence"] + div.chunk[data-level="sentence"] { + margin-top: -0.2em; } p { diff --git a/abogen/web/templates/reader_embed.html b/abogen/web/templates/reader_embed.html index 1cc21a4..48d65e8 100644 --- a/abogen/web/templates/reader_embed.html +++ b/abogen/web/templates/reader_embed.html @@ -126,11 +126,19 @@ } #reader .chunk { - margin-bottom: 1.75rem; + margin-bottom: 1.25rem; white-space: pre-wrap; word-wrap: break-word; } + #reader .chunk[data-level="sentence"] { + margin-bottom: 0.6rem; + } + + #reader .chunk[data-level="sentence"] + #reader .chunk[data-level="sentence"] { + margin-top: -0.2rem; + } + #reader p, #reader li, #reader blockquote { diff --git a/tests/test_chunk_helpers.py b/tests/test_chunk_helpers.py index 096bb7c..a677249 100644 --- a/tests/test_chunk_helpers.py +++ b/tests/test_chunk_helpers.py @@ -57,3 +57,23 @@ def test_chunk_text_merges_title_abbreviations() -> None: assert normalized_value assert text_value.startswith("Dr.") assert "Doctor" in normalized_value + display_value = str(chunk.get("display_text") or "") + assert display_value.startswith("Dr.") + + +def test_chunk_text_display_preserves_whitespace() -> None: + text = "Line one with double spaces.\nSecond line\n\nThird paragraph." + + chunks = chunk_text( + chapter_index=0, + chapter_title="Chapter 1", + text=text, + level="paragraph", + ) + + assert len(chunks) == 2 + first_display = str(chunks[0].get("display_text") or "") + assert " with " in first_display + assert first_display.endswith("\n\n") + second_display = str(chunks[1].get("display_text") or "") + assert second_display == "Third paragraph."