mirror of
https://github.com/denizsafak/abogen.git
synced 2026-07-18 05:40:26 +02:00
276 lines
8.9 KiB
Python
276 lines
8.9 KiB
Python
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass
|
||
from typing import Dict, Iterable, Iterator, List, Literal, Optional, Tuple
|
||
from typing import Pattern
|
||
|
||
import re
|
||
|
||
from abogen.kokoro_text_normalization import ApostropheConfig, normalize_for_pipeline
|
||
from abogen.normalization_settings import build_apostrophe_config, get_runtime_settings
|
||
|
||
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)
|
||
class Chunk:
|
||
id: str
|
||
chapter_index: int
|
||
chunk_index: int
|
||
level: ChunkLevel
|
||
text: str
|
||
speaker_id: str = "narrator"
|
||
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 {
|
||
"id": self.id,
|
||
"chapter_index": self.chapter_index,
|
||
"chunk_index": self.chunk_index,
|
||
"level": self.level,
|
||
"text": self.text,
|
||
"speaker_id": self.speaker_id,
|
||
"voice": self.voice,
|
||
"voice_profile": self.voice_profile,
|
||
"voice_formula": self.voice_formula,
|
||
"display_text": self.display_text,
|
||
}
|
||
|
||
|
||
def _iter_paragraphs(text: str) -> Iterator[str]:
|
||
for raw_segment in _PARAGRAPH_SPLIT_REGEX.split(text.strip()):
|
||
normalized = raw_segment.strip()
|
||
if normalized:
|
||
yield normalized
|
||
|
||
|
||
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()
|
||
raw_segment = paragraph[start:end]
|
||
candidate = raw_segment.strip()
|
||
if candidate:
|
||
yield candidate, raw_segment
|
||
start = match.end()
|
||
tail_raw = paragraph[start:]
|
||
tail = tail_raw.strip()
|
||
if tail:
|
||
yield tail, tail_raw
|
||
|
||
|
||
def _normalize_whitespace(value: str) -> str:
|
||
return _WHITESPACE_REGEX.sub(" ", value).strip()
|
||
|
||
|
||
def _normalize_chunk_text(value: str) -> str:
|
||
settings = get_runtime_settings()
|
||
config = build_apostrophe_config(
|
||
settings=settings, base=_PIPELINE_APOSTROPHE_CONFIG
|
||
)
|
||
normalized = normalize_for_pipeline(value, config=config, settings=settings)
|
||
return _normalize_whitespace(normalized)
|
||
|
||
|
||
def _split_sentences(paragraph: str) -> List[Tuple[str, str]]:
|
||
sentences = list(_iter_sentences(paragraph))
|
||
if not sentences:
|
||
return []
|
||
|
||
merged: List[Tuple[str, str]] = []
|
||
buffer_norm: List[str] = []
|
||
buffer_raw: List[str] = []
|
||
|
||
for normalized_sentence, raw_sentence in sentences:
|
||
if buffer_norm:
|
||
buffer_norm.append(normalized_sentence)
|
||
buffer_raw.append(raw_sentence)
|
||
else:
|
||
buffer_norm = [normalized_sentence]
|
||
buffer_raw = [raw_sentence]
|
||
|
||
if _ABBREVIATION_END_RE.search(normalized_sentence.rstrip()):
|
||
continue
|
||
|
||
merged.append((" ".join(buffer_norm), "".join(buffer_raw)))
|
||
buffer_norm = []
|
||
buffer_raw = []
|
||
|
||
if buffer_norm:
|
||
merged.append((" ".join(buffer_norm), "".join(buffer_raw)))
|
||
|
||
return merged
|
||
|
||
|
||
def chunk_text(
|
||
*,
|
||
chapter_index: int,
|
||
chapter_title: str,
|
||
text: str,
|
||
level: ChunkLevel,
|
||
speaker_id: str = "narrator",
|
||
voice: Optional[str] = None,
|
||
voice_profile: Optional[str] = None,
|
||
voice_formula: Optional[str] = None,
|
||
chunk_prefix: Optional[str] = None,
|
||
) -> List[Dict[str, object]]:
|
||
"""Split text into ordered chunk dictionaries."""
|
||
|
||
prefix = chunk_prefix or f"chap{chapter_index:04d}"
|
||
chunks: List[Dict[str, object]] = []
|
||
|
||
if level == "paragraph":
|
||
paragraphs = list(_iter_paragraphs(text)) or [text.strip()]
|
||
for para_index, paragraph in enumerate(paragraphs):
|
||
normalized = _normalize_whitespace(paragraph)
|
||
if not normalized:
|
||
continue
|
||
chunk_id = f"{prefix}_p{para_index:04d}"
|
||
payload = Chunk(
|
||
id=chunk_id,
|
||
chapter_index=chapter_index,
|
||
chunk_index=len(chunks),
|
||
level=level,
|
||
text=normalized,
|
||
speaker_id=speaker_id,
|
||
voice=voice,
|
||
voice_profile=voice_profile,
|
||
voice_formula=voice_formula,
|
||
).as_dict()
|
||
payload["normalized_text"] = _normalize_chunk_text(paragraph)
|
||
payload["original_text"] = paragraph
|
||
chunks.append(payload)
|
||
_attach_display_text(text, chunks)
|
||
return chunks
|
||
|
||
# Sentence level – flatten paragraphs into individual sentences
|
||
sentence_index = 0
|
||
for para_index, paragraph in enumerate(
|
||
list(_iter_paragraphs(text)) or [text.strip()]
|
||
):
|
||
normalized_para = _normalize_whitespace(paragraph)
|
||
if not normalized_para:
|
||
continue
|
||
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}"
|
||
payload = Chunk(
|
||
id=chunk_id,
|
||
chapter_index=chapter_index,
|
||
chunk_index=sentence_index,
|
||
level=level,
|
||
text=normalized_sentence,
|
||
speaker_id=speaker_id,
|
||
voice=voice,
|
||
voice_profile=voice_profile,
|
||
voice_formula=voice_formula,
|
||
).as_dict()
|
||
payload["normalized_text"] = _normalize_chunk_text(raw_sentence)
|
||
payload["display_text"] = raw_sentence
|
||
payload["original_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)
|
||
chunk.setdefault("original_text", chunk.get("display_text") or candidate)
|
||
continue
|
||
start, end = match
|
||
chunk["display_text"] = source[start:end]
|
||
chunk["original_text"] = source[start:end]
|
||
cursor = end
|
||
|
||
|
||
def build_chunks_for_chapters(
|
||
chapters: Iterable[Dict[str, object]],
|
||
*,
|
||
level: ChunkLevel,
|
||
speaker_id: str = "narrator",
|
||
) -> List[Dict[str, object]]:
|
||
"""Generate chunk dictionaries for a sequence of chapter payloads."""
|
||
all_chunks: List[Dict[str, object]] = []
|
||
for chapter_index, entry in enumerate(chapters):
|
||
if not isinstance(entry, dict): # defensive
|
||
continue
|
||
text = str(entry.get("text", "") or "").strip()
|
||
if not text:
|
||
continue
|
||
voice = entry.get("voice")
|
||
voice_profile = entry.get("voice_profile")
|
||
voice_formula = entry.get("voice_formula")
|
||
prefix = entry.get("id") or f"chap{chapter_index:04d}"
|
||
chapter_chunks = chunk_text(
|
||
chapter_index=chapter_index,
|
||
chapter_title=str(entry.get("title") or f"Chapter {chapter_index + 1}"),
|
||
text=text,
|
||
level=level,
|
||
speaker_id=speaker_id,
|
||
voice=str(voice) if voice else None,
|
||
voice_profile=str(voice_profile) if voice_profile else None,
|
||
voice_formula=str(voice_formula) if voice_formula else None,
|
||
chunk_prefix=str(prefix),
|
||
)
|
||
all_chunks.extend(chapter_chunks)
|
||
return all_chunks
|