feat: Enhance text normalization and chunking logic to preserve original whitespace and handle abbreviations

This commit is contained in:
JB
2025-10-10 08:38:00 -07:00
parent 258c3549f7
commit 3a91e79cb6
9 changed files with 251 additions and 71 deletions
+46 -7
View File
@@ -5,11 +5,19 @@ from typing import Dict, Iterable, Iterator, List, Literal, Optional
import re
from abogen.kokoro_text_normalization import ApostropheConfig, normalize_for_pipeline
ChunkLevel = Literal["paragraph", "sentence"]
_SENTENCE_SPLIT_REGEX = re.compile(r"(?<!\b[A-Z])[.!?][\s\n]+")
_WHITESPACE_REGEX = re.compile(r"\s+")
_PARAGRAPH_SPLIT_REGEX = re.compile(r"(?:\r?\n){2,}")
_ABBREVIATION_END_RE = re.compile(
r"\b(?:Mr|Mrs|Ms|Dr|Prof|Rev|Sr|Jr|St|Gen|Lt|Col|Sgt|Capt|Adm|Cmdr|vs|etc)\.$",
re.IGNORECASE,
)
_PIPELINE_APOSTROPHE_CONFIG = ApostropheConfig()
@dataclass(frozen=True)
@@ -64,6 +72,37 @@ def _normalize_whitespace(value: str) -> str:
return _WHITESPACE_REGEX.sub(" ", value).strip()
def _normalize_chunk_text(value: str) -> str:
normalized = normalize_for_pipeline(value, config=_PIPELINE_APOSTROPHE_CONFIG)
return _normalize_whitespace(normalized)
def _split_sentences(paragraph: str) -> List[str]:
sentences = list(_iter_sentences(paragraph))
if not sentences:
return []
merged: List[str] = []
buffer: List[str] = []
for sentence in sentences:
if buffer:
buffer.append(sentence)
else:
buffer = [sentence]
if _ABBREVIATION_END_RE.search(sentence.rstrip()):
continue
merged.append(" ".join(buffer))
buffer = []
if buffer:
merged.append(" ".join(buffer))
return merged
def chunk_text(
*,
chapter_index: int,
@@ -88,8 +127,7 @@ def chunk_text(
if not normalized:
continue
chunk_id = f"{prefix}_p{para_index:04d}"
chunks.append(
Chunk(
payload = Chunk(
id=chunk_id,
chapter_index=chapter_index,
chunk_index=len(chunks),
@@ -100,7 +138,8 @@ def chunk_text(
voice_profile=voice_profile,
voice_formula=voice_formula,
).as_dict()
)
payload["normalized_text"] = _normalize_chunk_text(paragraph)
chunks.append(payload)
return chunks
# Sentence level flatten paragraphs into individual sentences
@@ -109,14 +148,13 @@ def chunk_text(
normalized_para = _normalize_whitespace(paragraph)
if not normalized_para:
continue
sentences = list(_iter_sentences(normalized_para)) or [normalized_para]
sentences = _split_sentences(normalized_para) or [normalized_para]
for sent_local_index, sentence in enumerate(sentences):
normalized_sentence = _normalize_whitespace(sentence)
if not normalized_sentence:
continue
chunk_id = f"{prefix}_p{para_index:04d}_s{sent_local_index:04d}"
chunks.append(
Chunk(
payload = Chunk(
id=chunk_id,
chapter_index=chapter_index,
chunk_index=sentence_index,
@@ -127,7 +165,8 @@ def chunk_text(
voice_profile=voice_profile,
voice_formula=voice_formula,
).as_dict()
)
payload["normalized_text"] = _normalize_chunk_text(sentence)
chunks.append(payload)
sentence_index += 1
return chunks
+61 -27
View File
@@ -1,13 +1,14 @@
from __future__ import annotations
import html
import re
import shutil
import uuid
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Any, Dict, Iterable, List, Optional, Sequence
from typing import Any, Dict, Iterable, List, Optional, Pattern, Sequence, Tuple
import zipfile
from abogen.text_extractor import ExtractedChapter, ExtractionResult
@@ -259,6 +260,13 @@ class EPUB3PackageBuilder:
)
)
chapter_text = ""
if 0 <= chapter_index < len(self.extraction.chapters):
chapter_entry = self.extraction.chapters[chapter_index]
chapter_text = getattr(chapter_entry, "text", "") or ""
_restore_original_chunk_text(chapter_text, overlays)
return overlays
def _render_chapter_xhtml(self, chapter: ChapterDocument) -> str:
@@ -617,35 +625,18 @@ 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 ""
paragraphs = _split_paragraphs(chunk.text)
if not paragraphs:
paragraphs = ["&nbsp;"]
return " <div class=\"chunk\" id=\"{id}\"{speaker}{voice}>\n{body}\n </div>".format(
id=escaped_id,
speaker=speaker_attr,
voice=voice_attr,
body="\n".join(f" <p>{para}</p>" for para in paragraphs),
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" {body}\n"
" </div>"
)
def _split_paragraphs(text: str) -> List[str]:
if not text:
return []
segments = [segment.strip() for segment in text.replace("\r", "").split("\n\n")]
paragraphs: List[str] = []
for segment in segments:
if not segment:
continue
lines = [html.escape(line.strip()) for line in segment.split("\n") if line.strip()]
if not lines:
continue
if len(lines) == 1:
paragraphs.append(lines[0])
else:
paragraphs.append("<br />".join(lines))
return paragraphs
def _format_smil_time(value: Optional[float]) -> str:
if value is None or value < 0:
value = 0.0
@@ -672,6 +663,49 @@ def _safe_float(value: Any) -> Optional[float]:
return None
def _restore_original_chunk_text(chapter_text: str, overlays: List[ChunkOverlay]) -> None:
if not chapter_text or not overlays:
return
cursor = 0
for chunk in overlays:
candidate = chunk.text or ""
if not candidate:
continue
match = _search_original_span(chapter_text, candidate, cursor)
if match is None and cursor:
match = _search_original_span(chapter_text, candidate, 0)
if match is None:
continue
start, end = match
chunk.text = chapter_text[start:end]
cursor = end
def _search_original_span(source: str, normalized: str, start: int) -> Optional[Tuple[int, int]]:
if not normalized:
return None
pattern = _build_chunk_pattern(normalized)
match = pattern.search(source, start)
if not match:
return None
return match.start(1), match.end(1)
_CHUNK_REGEX_CACHE: Dict[str, Pattern[str]] = {}
def _build_chunk_pattern(text: str) -> Pattern[str]:
cached = _CHUNK_REGEX_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)
_CHUNK_REGEX_CACHE[text] = pattern
return pattern
def _render_metadata_xml(
title: str,
authors: Sequence[str],
+21 -1
View File
@@ -2,7 +2,7 @@ from __future__ import annotations
import re
import unicodedata
from dataclasses import dataclass
from typing import List, Tuple, Iterable, Callable
from typing import List, Tuple, Iterable, Callable, Optional
# ---------- Configuration Dataclass ----------
@@ -416,6 +416,26 @@ def apply_phoneme_hints(text: str, iz_marker="IZ") -> str:
"""
return text.replace(iz_marker, " iz")
DEFAULT_APOSTROPHE_CONFIG = ApostropheConfig()
def normalize_for_pipeline(text: str, *, config: Optional[ApostropheConfig] = None) -> str:
"""Normalize text for the synthesis pipeline.
This expands contractions, normalizes apostrophes, and adds phoneme hints
using the provided configuration so downstream chunking and synthesis share
the same representation.
"""
cfg = config or DEFAULT_APOSTROPHE_CONFIG
normalized, _details = normalize_apostrophes(text, cfg)
normalized = expand_titles_and_suffixes(normalized)
normalized = ensure_terminal_punctuation(normalized)
if cfg.add_phoneme_hints:
normalized = apply_phoneme_hints(normalized, iz_marker=cfg.sibilant_iz_marker)
return normalized
# ---------- Example Usage ----------
if __name__ == "__main__":
+1 -1
View File
@@ -175,7 +175,7 @@ def analyze_speakers(
for chunk in ordered_chunks:
chunk_id = str(chunk.get("id") or "")
text = str(chunk.get("text") or "")
text = str(chunk.get("normalized_text") or chunk.get("text") or "")
speaker_id, confidence, quote = _infer_chunk_speaker(text, last_explicit)
if speaker_id is None:
speaker_id = last_explicit or narrator_id
+7 -14
View File
@@ -20,13 +20,7 @@ import static_ffmpeg
from abogen.constants import VOICES_INTERNAL
from abogen.epub3.exporter import build_epub3_package
from abogen.kokoro_text_normalization import (
ApostropheConfig,
apply_phoneme_hints,
expand_titles_and_suffixes,
ensure_terminal_punctuation,
normalize_apostrophes,
)
from abogen.kokoro_text_normalization import ApostropheConfig, normalize_for_pipeline
from abogen.text_extractor import ExtractedChapter, extract_from_path
from abogen.utils import (
calculate_text_length,
@@ -402,12 +396,7 @@ _APOSTROPHE_CONFIG = ApostropheConfig()
def _normalize_for_pipeline(text: str) -> str:
normalized, _details = normalize_apostrophes(text, _APOSTROPHE_CONFIG)
normalized = expand_titles_and_suffixes(normalized)
normalized = ensure_terminal_punctuation(normalized)
if _APOSTROPHE_CONFIG.add_phoneme_hints:
return apply_phoneme_hints(normalized, iz_marker=_APOSTROPHE_CONFIG.sibilant_iz_marker)
return normalized
return normalize_for_pipeline(text, config=_APOSTROPHE_CONFIG)
def _chapter_voice_spec(job: Job, override: Optional[Dict[str, Any]]) -> str:
@@ -988,7 +977,11 @@ def run_conversion_job(job: Job) -> None:
level="debug",
)
for chunk_entry in chunks_for_chapter:
chunk_text = str(chunk_entry.get("text") or "").strip()
chunk_text = str(
chunk_entry.get("normalized_text")
or chunk_entry.get("text")
or ""
).strip()
if not chunk_text:
continue
+4
View File
@@ -1084,6 +1084,10 @@ class ConversionService:
else:
chunk["text"] = ""
normalized_value = entry.get("normalized_text")
if normalized_value is not None:
chunk["normalized_text"] = str(normalized_value)
speaker_value = entry.get("speaker_id", entry.get("speaker"))
chunk["speaker_id"] = str(speaker_value) if speaker_value else "narrator"
+6
View File
@@ -125,6 +125,12 @@
outline-offset: 4px;
}
#reader .chunk {
margin-bottom: 1.75rem;
white-space: pre-wrap;
word-wrap: break-word;
}
#reader p,
#reader li,
#reader blockquote {
+20
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
from types import SimpleNamespace
from abogen.chunking import chunk_text
from abogen.web.conversion_runner import _chunk_voice_spec, _group_chunks_by_chapter
@@ -37,3 +38,22 @@ def test_chunk_voice_spec_uses_fallback_when_no_overrides() -> None:
chunk = {"speaker_id": "unknown"}
assert _chunk_voice_spec(job, chunk, "fallback") == "fallback"
def test_chunk_text_merges_title_abbreviations() -> None:
text = "Dr. Watson met Mr. Holmes at 5 p.m."
chunks = chunk_text(
chapter_index=0,
chapter_title="Chapter 1",
text=text,
level="sentence",
)
assert len(chunks) == 1
chunk = chunks[0]
text_value = str(chunk["text"])
normalized_value = str(chunk.get("normalized_text") or "")
assert normalized_value
assert text_value.startswith("Dr.")
assert "Doctor" in normalized_value
+64
View File
@@ -105,3 +105,67 @@ def test_build_epub3_package_handles_missing_markers(tmp_path) -> None:
assert "Chapter 1" in nav_doc
chapter_doc = archive.read("OEBPS/text/chapter_0001.xhtml").decode("utf-8")
assert "Hello world." in chapter_doc
def test_epub3_preserves_original_whitespace(tmp_path) -> None:
extraction = ExtractionResult(
chapters=[
ExtractedChapter(
title="Intro",
text="Line one with double spaces.\nSecond line\n\nThird paragraph.",
)
],
metadata={"title": "Sample", "artist": "Author", "language": "en"},
)
chunks = [
{
"id": "chap0000_p0000",
"chapter_index": 0,
"chunk_index": 0,
"text": "Line one with double spaces.",
"speaker_id": "narrator",
},
{
"id": "chap0000_p0001",
"chapter_index": 0,
"chunk_index": 1,
"text": "Second line",
"speaker_id": "narrator",
},
{
"id": "chap0000_p0002",
"chapter_index": 0,
"chunk_index": 2,
"text": "Third paragraph.",
"speaker_id": "narrator",
},
]
chunk_markers = [
{"id": chunk["id"], "chapter_index": 0, "chunk_index": chunk["chunk_index"], "start": None, "end": None}
for chunk in chunks
]
metadata_tags = {"title": "Sample", "artist": "Author", "language": "en"}
audio_path = tmp_path / "audio.mp3"
audio_path.write_bytes(b"ID3 audio")
output_path = tmp_path / "output.epub"
build_epub3_package(
output_path=output_path,
book_id="job-whitespace",
extraction=extraction,
metadata_tags=metadata_tags,
chapter_markers=[],
chunk_markers=chunk_markers,
chunks=chunks,
audio_path=audio_path,
speaker_mode="single",
)
with zipfile.ZipFile(output_path) as archive:
chapter_doc = archive.read("OEBPS/text/chapter_0001.xhtml").decode("utf-8")
assert "Line one with double spaces." in chapter_doc
normalized = chapter_doc.replace(" ", "")
assert "Second line\n\nThird paragraph." in normalized