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

This commit is contained in:
JB
2025-10-10 11:17:14 -07:00
parent 443dee09b6
commit 15c1220ba3
4 changed files with 144 additions and 27 deletions
+77 -22
View File
@@ -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]],
*,
+38 -4
View File
@@ -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 {
+9 -1
View File
@@ -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 {