feat: Enhance chunking logic to include display text and preserve whitespace in sentences

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] = []
+    merged: List[Tuple[str, str]] = []
-    buffer: List[str] = []
+    buffer_norm: List[str] = []
+    buffer_raw: List[str] = []
 
-    for sentence in sentences:
+    for normalized_sentence, raw_sentence in sentences:
-        if buffer:
+        if buffer_norm:
-            buffer.append(sentence)
+            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))
+        merged.append((" ".join(buffer_norm), "".join(buffer_raw)))
-        buffer = []
+        buffer_norm = []
+        buffer_raw = []
 
-    if buffer:
+    if buffer_norm:
-        merged.append(" ".join(buffer))
+        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]
+        sentence_pairs = _split_sentences(paragraph) or [(normalized_para, paragraph)]
-        for sent_local_index, sentence in enumerate(sentences):
+        for sent_local_index, (normalized_sentence, raw_sentence) in enumerate(sentence_pairs):
-            normalized_sentence = _normalize_whitespace(sentence)
+            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]],
     *,

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 = "<p></p>"
+        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 = (
+                    "      <pre class=\"chapter-original\" hidden>\n"
+                    f"{safe_original}\n"
+                    "      </pre>"
+                )
 
         return (
             "<?xml version=\"1.0\" encoding=\"utf-8\"?>\n"
@@ -288,10 +305,17 @@ class EPUB3PackageBuilder:
             "    <section epub:type=\"chapter\" id=\"chapter-{index:04d}\">\n"
             "      <h1>{title}</h1>\n"
             "      {chunks}\n"
+            "{original_block}"
             "    </section>\n"
             "  </body>\n"
             "</html>\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 = "&nbsp;"
     body = escaped_text.replace("\n", "\n        ")
     return (
-        f"      <div class=\"chunk\" id=\"{escaped_id}\"{speaker_attr}{voice_attr}>\n"
+        f"      <div class=\"chunk\" id=\"{escaped_id}\"{speaker_attr}{voice_attr}{level_attr}>\n"
         f"        {body}\n"
         "      </div>"
     )
@@ -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 {

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 {

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."