From 2c61f55f81ae481cd85435234ced8a29fcdb34ea Mon Sep 17 00:00:00 2001 From: Artem Akymenko Date: Sun, 26 Jul 2026 06:49:34 +0000 Subject: [PATCH] refactor: spaCy pre-TTS segmentation moved to shared layer - conversion_pipeline.py: new spacy_pre_tts_segmentation() function - Handles: condition checks, language exclusion, split_pattern override - Executor: integrated spaCy pre-TTS before each synthesize_text() call - 6 new tests for condition checks and fallback behavior - English excluded from pre-TTS (spaCy only for post-TTS subtitles) - Fallback to regex when spaCy unavailable --- abogen/application/conversion_executor.py | 39 ++++++-- abogen/domain/conversion_pipeline.py | 110 +++++++++++++++++++++- tests/test_domain_conversion_pipeline.py | 80 ++++++++++++++++ 3 files changed, 218 insertions(+), 11 deletions(-) diff --git a/abogen/application/conversion_executor.py b/abogen/application/conversion_executor.py index 8242ace..e13adc8 100644 --- a/abogen/application/conversion_executor.py +++ b/abogen/application/conversion_executor.py @@ -425,16 +425,37 @@ def execute_conversion( # Track voice for chapter marker collector.on_segment(seg_provider, seg_voice, segment.voice_spec) - seg_start_time = stats.current_time - local_segments, accumulated_tokens = synthesize_text( - text=seg_text, - params=synth, - backend=seg_backend, - voice=seg_voice, - speed=seg_speed or request.speed, - chapter_sink=chapter_sink, - preview_callback=lambda text: events.log(f" {text[:80]}"), + # spaCy pre-TTS segmentation + from abogen.domain.conversion_pipeline import spacy_pre_tts_segmentation + + is_subtitle_input = bool( + request.subtitle_input ) + spacy_segments, active_split = spacy_pre_tts_segmentation( + seg_text, + request.language, + request.subtitle_mode, + is_subtitle_input=is_subtitle_input, + use_spacy_segmentation=use_spacy, + log_callback=lambda msg: events.log(msg), + ) + + seg_start_time = stats.current_time + accumulated_tokens: List[Dict[str, Any]] = [] + for spacy_seg in spacy_segments: + if not spacy_seg.strip(): + continue + _, seg_tokens = synthesize_text( + text=spacy_seg, + params=synth, + backend=seg_backend, + voice=seg_voice, + speed=seg_speed or request.speed, + chapter_sink=chapter_sink, + preview_callback=lambda text: events.log(f" {text[:80]}"), + split_pattern_override=active_split, + ) + accumulated_tokens.extend(seg_tokens) # Process subtitles if audio_sink and accumulated_tokens: diff --git a/abogen/domain/conversion_pipeline.py b/abogen/domain/conversion_pipeline.py index 6c0c023..fe3610b 100644 --- a/abogen/domain/conversion_pipeline.py +++ b/abogen/domain/conversion_pipeline.py @@ -9,8 +9,8 @@ from __future__ import annotations import logging from dataclasses import dataclass, field -from abogen.domain.enums import SubtitleMode -from typing import Any, Dict, Iterator, List, Optional +from abogen.domain.enums import Language, SubtitleMode +from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple import numpy as np @@ -21,6 +21,112 @@ from abogen.domain.audio_buffer import SAMPLE_RATE logger = logging.getLogger(__name__) +# Languages where spaCy is used for pre-TTS segmentation +# English ("a", "b") is excluded — spaCy only used for post-TTS subtitles +_SPACY_EXCLUDED_LANGS = {Language.EN_US, Language.EN_GB} + +# CJK languages — different spacing pattern +_CJK_LANGS = {Language.ZH, Language.JA} + + +def spacy_pre_tts_segmentation( + text: str, + lang_code: Any, + subtitle_mode: Any, + *, + is_subtitle_input: bool = False, + use_spacy_segmentation: bool = True, + log_callback: Optional[Callable[[str], None]] = None, +) -> Tuple[List[str], str]: + """Segment text using spaCy before TTS, with split_pattern override. + + For non-English languages, spaCy sentence segmentation produces better + sentence boundaries than regex. This function: + 1. Checks if spaCy should be used (toggle on, not disabled mode, not subtitle input) + 2. For non-English: runs spaCy segmentation, computes split_pattern override + 3. For English: returns single segment with default pattern (spaCy only for subtitles) + 4. If spaCy fails: falls back to default pattern + + Args: + text: Text to segment. + lang_code: Language code (Language enum or string like "a", "de", "fr"). + subtitle_mode: SubtitleMode enum or string. + is_subtitle_input: True if source is .srt/.ass/.vtt file. + use_spacy_segmentation: User toggle for spaCy segmentation. + log_callback: Optional logging function. + + Returns: + Tuple of (text_segments, active_split_pattern). + text_segments is a list of sentences (always at least one element). + active_split_pattern is the regex to use for TTS backend splitting. + """ + from abogen.domain.split_pattern import PUNCTUATION_COMMAS, get_split_pattern + + def _log(msg: str) -> None: + if log_callback: + log_callback(msg) + + # Normalize language + lang_enum = _to_language_enum(lang_code) + + # Default split pattern + default_split = get_split_pattern(lang_code, subtitle_mode) + + # Check conditions + if not use_spacy_segmentation: + return [text], default_split + + subtitle_mode_str = _to_subtitle_mode_str(subtitle_mode) + if subtitle_mode_str in ("Disabled", "Line"): + return [text], default_split + + if is_subtitle_input: + return [text], default_split + + # English: spaCy only for post-TTS subtitles, not pre-TTS + if lang_enum in _SPACY_EXCLUDED_LANGS: + return [text], default_split + + # Non-English: run spaCy pre-TTS segmentation + from abogen.spacy_utils import segment_sentences + + _log("Using spaCy for sentence segmentation (pre-TTS)...") + spacy_sentences = segment_sentences(text, lang_code, log_callback=log_callback) + + if not spacy_sentences: + _log("spaCy: Fallback to default segmentation...") + return [text], default_split + + _log(f"spaCy: Text segmented into {len(spacy_sentences)} sentences...") + + # Compute split_pattern override based on subtitle mode + spacing_pattern = r"\s*" if lang_enum in _CJK_LANGS else r"\s+" + + if subtitle_mode_str == "Sentence + Comma": + active_split = r"(?<=[{}]){}|\n+".format(PUNCTUATION_COMMAS, spacing_pattern) + else: + # Sentence mode: spaCy already split, only split on newlines + active_split = "\n" + + return spacy_sentences, active_split + + +def _to_language_enum(lang_code: Any) -> Language: + """Convert lang_code to Language enum.""" + if isinstance(lang_code, Language): + return lang_code + try: + return Language.from_str(str(lang_code)) + except (ValueError, AttributeError): + return Language.EN_US + + +def _to_subtitle_mode_str(subtitle_mode: Any) -> str: + """Convert subtitle_mode to string.""" + if isinstance(subtitle_mode, SubtitleMode): + return subtitle_mode.value + return str(subtitle_mode) + @dataclass class SegmentResult: diff --git a/tests/test_domain_conversion_pipeline.py b/tests/test_domain_conversion_pipeline.py index c455a6c..5970e14 100644 --- a/tests/test_domain_conversion_pipeline.py +++ b/tests/test_domain_conversion_pipeline.py @@ -181,3 +181,83 @@ class TestTtsSegments: )) assert results[0].chunk_start == 10.0 assert results[1].chunk_start == 11.0 + + +class TestSpacyPreTtsSegmentation: + """Tests for spacy_pre_tts_segmentation().""" + + def test_disabled_when_use_spacy_false(self): + from abogen.domain.conversion_pipeline import spacy_pre_tts_segmentation + from abogen.domain.enums import Language + + segments, split = spacy_pre_tts_segmentation( + "Hello world", + Language.EN_US, + "Disabled", + use_spacy_segmentation=False, + ) + assert segments == ["Hello world"] + assert isinstance(split, str) + + def test_disabled_for_disabled_subtitle_mode(self): + from abogen.domain.conversion_pipeline import spacy_pre_tts_segmentation + from abogen.domain.enums import Language + + segments, split = spacy_pre_tts_segmentation( + "Hello world", + Language.FR, + "Disabled", + use_spacy_segmentation=True, + ) + assert segments == ["Hello world"] + + def test_disabled_for_line_subtitle_mode(self): + from abogen.domain.conversion_pipeline import spacy_pre_tts_segmentation + from abogen.domain.enums import Language + + segments, split = spacy_pre_tts_segmentation( + "Hello world", + Language.ES, + "Line", + use_spacy_segmentation=True, + ) + assert segments == ["Hello world"] + + def test_disabled_for_subtitle_input(self): + from abogen.domain.conversion_pipeline import spacy_pre_tts_segmentation + from abogen.domain.enums import Language + + segments, split = spacy_pre_tts_segmentation( + "Hello world", + Language.FR, + "Sentence", + is_subtitle_input=True, + use_spacy_segmentation=True, + ) + assert segments == ["Hello world"] + + def test_english_excluded_from_pre_tts(self): + """English uses spaCy only for post-TTS subtitles, not pre-TTS.""" + from abogen.domain.conversion_pipeline import spacy_pre_tts_segmentation + from abogen.domain.enums import Language + + segments, split = spacy_pre_tts_segmentation( + "Hello world. How are you?", + Language.EN_US, + "Sentence", + use_spacy_segmentation=True, + ) + # English should return single segment (no pre-TTS segmentation) + assert len(segments) == 1 + + def test_returns_at_least_one_segment(self): + from abogen.domain.conversion_pipeline import spacy_pre_tts_segmentation + from abogen.domain.enums import Language + + segments, split = spacy_pre_tts_segmentation( + "", + Language.FR, + "Sentence", + use_spacy_segmentation=True, + ) + assert len(segments) >= 1