mirror of
https://github.com/denizsafak/abogen.git
synced 2026-07-18 13:40:27 +02:00
refactor: extract subtitle token processing to domain layer
- Add abogen/domain/subtitle_generation.py with: - process_subtitle_tokens(): main function for converting TTS tokens to subtitles - Support for all subtitle modes: Line, Sentence, Sentence + Comma, Sentence + Highlighting - Support for word-count based grouping (e.g., '5' for 5 words per entry) - spaCy integration for English sentence boundary detection - Karaoke highlighting tags for Sentence + Highlighting mode - Punctuation constants for sentence splitting - Update abogen/pyqt/conversion.py: - Replace _process_subtitle_tokens method body with call to domain function - Remove ~260 lines of duplicate logic - Add tests/test_subtitle_generation.py with comprehensive unit tests
This commit is contained in:
Artem Akymenko
@@ -0,0 +1,358 @@
|
||||
"""Subtitle generation utilities for audiobook generation.
|
||||
|
||||
This module provides functions for processing TTS tokens into subtitle entries
|
||||
according to various subtitle modes (Line, Sentence, Sentence + Comma,
|
||||
Sentence + Highlighting).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
|
||||
# Punctuation constants for sentence splitting
|
||||
PUNCTUATION_SENTENCE = ".!?\u061f\u3002\uff01\uff1f" # .!? .?. ??
|
||||
PUNCTUATION_SENTENCE_COMMA = ".!?,\u3001\u061f\u3002\uff01\uff0c\uff1f" # .!?, ,. ??
|
||||
|
||||
|
||||
def process_subtitle_tokens(
|
||||
tokens_with_timestamps: List[dict],
|
||||
subtitle_entries: List[Tuple[float, float, str]],
|
||||
max_subtitle_words: int,
|
||||
subtitle_mode: str,
|
||||
lang_code: str,
|
||||
use_spacy_segmentation: bool = False,
|
||||
fallback_end_time: Optional[float] = None,
|
||||
) -> None:
|
||||
"""Process TTS tokens into subtitle entries according to the subtitle mode.
|
||||
|
||||
This function modifies subtitle_entries in-place by appending new entries.
|
||||
|
||||
Args:
|
||||
tokens_with_timestamps: List of token dictionaries with 'start', 'end', 'text',
|
||||
and 'whitespace' keys.
|
||||
subtitle_entries: List to append subtitle entries to (modified in-place).
|
||||
Each entry is a tuple of (start_time, end_time, text).
|
||||
max_subtitle_words: Maximum number of words per subtitle entry.
|
||||
subtitle_mode: One of "Disabled", "Line", "Sentence", "Sentence + Comma",
|
||||
"Sentence + Highlighting", or a string like "5" for word-count mode.
|
||||
lang_code: Language code for spaCy processing (e.g., "a" for English).
|
||||
use_spacy_segmentation: Whether to use spaCy for sentence boundary detection.
|
||||
fallback_end_time: Fallback end time for the last entry if none is available.
|
||||
"""
|
||||
if not tokens_with_timestamps:
|
||||
return
|
||||
|
||||
processed_tokens = tokens_with_timestamps
|
||||
|
||||
# For English with spaCy enabled and sentence-based modes, use spaCy for sentence boundaries
|
||||
# spaCy is disabled when subtitle mode is "Disabled" or "Line"
|
||||
use_spacy_for_english = (
|
||||
use_spacy_segmentation
|
||||
and subtitle_mode not in ["Disabled", "Line"]
|
||||
and lang_code in ["a", "b"]
|
||||
and subtitle_mode in ["Sentence", "Sentence + Comma"]
|
||||
)
|
||||
|
||||
if subtitle_mode == "Sentence + Highlighting":
|
||||
_process_karaoke_highlighting(
|
||||
processed_tokens, subtitle_entries, max_subtitle_words, fallback_end_time
|
||||
)
|
||||
elif subtitle_mode in ["Sentence", "Sentence + Comma", "Line"]:
|
||||
if use_spacy_for_english and subtitle_mode != "Line":
|
||||
_process_spacy_sentences(
|
||||
processed_tokens, subtitle_entries, max_subtitle_words,
|
||||
subtitle_mode, lang_code, fallback_end_time
|
||||
)
|
||||
else:
|
||||
_process_regex_sentences(
|
||||
processed_tokens, subtitle_entries, max_subtitle_words,
|
||||
subtitle_mode, fallback_end_time
|
||||
)
|
||||
else:
|
||||
# Word count-based grouping (e.g., "5" for 5-word groups)
|
||||
_process_word_count(
|
||||
processed_tokens, subtitle_entries, max_subtitle_words,
|
||||
subtitle_mode, fallback_end_time
|
||||
)
|
||||
|
||||
|
||||
def _process_karaoke_highlighting(
|
||||
tokens: List[dict],
|
||||
subtitle_entries: List[Tuple[float, float, str]],
|
||||
max_subtitle_words: int,
|
||||
fallback_end_time: Optional[float],
|
||||
) -> None:
|
||||
"""Process tokens for Sentence + Highlighting mode (karaoke effect)."""
|
||||
separator = rf"[{re.escape(PUNCTUATION_SENTENCE)}]"
|
||||
current_sentence = []
|
||||
word_count = 0
|
||||
|
||||
for token in tokens:
|
||||
current_sentence.append(token)
|
||||
word_count += 1
|
||||
|
||||
# Split sentences based on separator or word count
|
||||
if (
|
||||
re.search(separator, token["text"]) and token.get("whitespace") == " "
|
||||
) or word_count >= max_subtitle_words:
|
||||
if current_sentence:
|
||||
# Create karaoke subtitle entry for this sentence
|
||||
start_time = current_sentence[0]["start"]
|
||||
end_time = current_sentence[-1]["end"]
|
||||
|
||||
# Generate karaoke text with timing
|
||||
karaoke_text = ""
|
||||
for t in current_sentence:
|
||||
# Calculate duration in centiseconds
|
||||
duration = (
|
||||
t["end"] - t["start"]
|
||||
if t.get("end") is not None and t.get("start") is not None
|
||||
else 0.5
|
||||
)
|
||||
duration_cs = int(duration * 100)
|
||||
# Add karaoke effect
|
||||
karaoke_text += f"{{\\kf{duration_cs}}}{t['text']}{t.get('whitespace', '') or ''}"
|
||||
|
||||
subtitle_entries.append(
|
||||
(start_time, end_time, karaoke_text.strip())
|
||||
)
|
||||
current_sentence = []
|
||||
word_count = 0
|
||||
|
||||
# Add any remaining tokens as a sentence
|
||||
if current_sentence:
|
||||
start_time = current_sentence[0]["start"]
|
||||
end_time = current_sentence[-1]["end"]
|
||||
|
||||
# Generate karaoke text for remaining tokens
|
||||
karaoke_text = ""
|
||||
for t in current_sentence:
|
||||
duration = t["end"] - t["start"] if t.get("end") and t.get("start") else 0.5
|
||||
duration_cs = int(duration * 100)
|
||||
karaoke_text += f"{{\\kf{duration_cs}}}{t['text']}{t.get('whitespace', '') or ''}"
|
||||
subtitle_entries.append((start_time, end_time, karaoke_text.strip()))
|
||||
|
||||
# Fallback for last entry
|
||||
_apply_fallback_end_time(subtitle_entries, fallback_end_time)
|
||||
|
||||
|
||||
def _process_spacy_sentences(
|
||||
tokens: List[dict],
|
||||
subtitle_entries: List[Tuple[float, float, str]],
|
||||
max_subtitle_words: int,
|
||||
subtitle_mode: str,
|
||||
lang_code: str,
|
||||
fallback_end_time: Optional[float],
|
||||
) -> None:
|
||||
"""Process tokens using spaCy for sentence boundary detection."""
|
||||
try:
|
||||
from abogen.spacy_utils import get_spacy_model
|
||||
except ImportError:
|
||||
# Fall back to regex if spaCy is not available
|
||||
_process_regex_sentences(
|
||||
tokens, subtitle_entries, max_subtitle_words,
|
||||
subtitle_mode, fallback_end_time
|
||||
)
|
||||
return
|
||||
|
||||
nlp = get_spacy_model(lang_code)
|
||||
if not nlp:
|
||||
_process_regex_sentences(
|
||||
tokens, subtitle_entries, max_subtitle_words,
|
||||
subtitle_mode, fallback_end_time
|
||||
)
|
||||
return
|
||||
|
||||
# Build full text and track character positions to token indices
|
||||
full_text = ""
|
||||
for token in tokens:
|
||||
text_part = token["text"] + (token.get("whitespace") or "")
|
||||
full_text += text_part
|
||||
|
||||
# Get sentence boundaries from spaCy
|
||||
doc = nlp(full_text)
|
||||
sentence_boundaries = [sent.end_char for sent in doc.sents]
|
||||
|
||||
# For "Sentence + Comma" mode, also split on commas
|
||||
if subtitle_mode == "Sentence + Comma":
|
||||
comma_positions = [
|
||||
i + 1 for i, c in enumerate(full_text) if c == ","
|
||||
]
|
||||
sentence_boundaries = sorted(
|
||||
set(sentence_boundaries + comma_positions)
|
||||
)
|
||||
|
||||
# Group tokens by sentence boundaries
|
||||
current_sentence = []
|
||||
word_count = 0
|
||||
current_char_pos = 0
|
||||
boundary_idx = 0
|
||||
|
||||
for token in tokens:
|
||||
current_sentence.append(token)
|
||||
word_count += 1
|
||||
text_len = len(token["text"]) + len(token.get("whitespace") or "")
|
||||
current_char_pos += text_len
|
||||
|
||||
# Check if we've hit a sentence boundary or max words
|
||||
at_boundary = (
|
||||
boundary_idx < len(sentence_boundaries)
|
||||
and current_char_pos >= sentence_boundaries[boundary_idx]
|
||||
)
|
||||
if at_boundary or word_count >= max_subtitle_words:
|
||||
if current_sentence:
|
||||
start_time = current_sentence[0]["start"]
|
||||
end_time = current_sentence[-1]["end"]
|
||||
sentence_text = "".join(
|
||||
t["text"] + (t.get("whitespace") or "")
|
||||
for t in current_sentence
|
||||
)
|
||||
subtitle_entries.append(
|
||||
(start_time, end_time, sentence_text.strip())
|
||||
)
|
||||
current_sentence = []
|
||||
word_count = 0
|
||||
if at_boundary:
|
||||
boundary_idx += 1
|
||||
|
||||
# Add remaining tokens
|
||||
if current_sentence:
|
||||
start_time = current_sentence[0]["start"]
|
||||
end_time = current_sentence[-1]["end"]
|
||||
sentence_text = "".join(
|
||||
t["text"] + (t.get("whitespace") or "")
|
||||
for t in current_sentence
|
||||
)
|
||||
subtitle_entries.append(
|
||||
(start_time, end_time, sentence_text.strip())
|
||||
)
|
||||
|
||||
# Fallback for last entry
|
||||
_apply_fallback_end_time(subtitle_entries, fallback_end_time)
|
||||
|
||||
|
||||
def _process_regex_sentences(
|
||||
tokens: List[dict],
|
||||
subtitle_entries: List[Tuple[float, float, str]],
|
||||
max_subtitle_words: int,
|
||||
subtitle_mode: str,
|
||||
fallback_end_time: Optional[float],
|
||||
) -> None:
|
||||
"""Process tokens using regex for sentence boundary detection."""
|
||||
# Define separator pattern based on mode
|
||||
if subtitle_mode == "Line":
|
||||
separator = r"\n"
|
||||
elif subtitle_mode == "Sentence":
|
||||
# Use punctuation without comma
|
||||
separator = rf"[{re.escape(PUNCTUATION_SENTENCE)}]"
|
||||
else: # Sentence + Comma
|
||||
# Use punctuation with comma
|
||||
separator = rf"[{re.escape(PUNCTUATION_SENTENCE_COMMA)}]"
|
||||
|
||||
current_sentence = []
|
||||
word_count = 0
|
||||
|
||||
for token in tokens:
|
||||
current_sentence.append(token)
|
||||
word_count += 1
|
||||
|
||||
# Split sentences based on separator or word count
|
||||
if (
|
||||
re.search(separator, token["text"]) and token.get("whitespace") == " "
|
||||
) or word_count >= max_subtitle_words:
|
||||
if current_sentence:
|
||||
# Create subtitle entry for this sentence
|
||||
start_time = current_sentence[0]["start"]
|
||||
end_time = current_sentence[-1]["end"]
|
||||
|
||||
# Simplified text joining logic
|
||||
sentence_text = ""
|
||||
for t in current_sentence:
|
||||
sentence_text += t["text"] + (t.get("whitespace") or "")
|
||||
|
||||
subtitle_entries.append(
|
||||
(start_time, end_time, sentence_text.strip())
|
||||
)
|
||||
current_sentence = []
|
||||
word_count = 0
|
||||
|
||||
# Add any remaining tokens as a sentence
|
||||
if current_sentence:
|
||||
start_time = current_sentence[0]["start"]
|
||||
end_time = current_sentence[-1]["end"]
|
||||
|
||||
# Simplified text joining logic
|
||||
sentence_text = ""
|
||||
for t in current_sentence:
|
||||
sentence_text += t["text"] + (t.get("whitespace") or "")
|
||||
subtitle_entries.append((start_time, end_time, sentence_text.strip()))
|
||||
|
||||
# Fallback for last entry
|
||||
_apply_fallback_end_time(subtitle_entries, fallback_end_time)
|
||||
|
||||
|
||||
def _process_word_count(
|
||||
tokens: List[dict],
|
||||
subtitle_entries: List[Tuple[float, float, str]],
|
||||
max_subtitle_words: int,
|
||||
subtitle_mode: str,
|
||||
fallback_end_time: Optional[float],
|
||||
) -> None:
|
||||
"""Process tokens by counting spaces (word count mode)."""
|
||||
try:
|
||||
word_count = int(subtitle_mode.split()[0])
|
||||
word_count = min(word_count, max_subtitle_words)
|
||||
except (ValueError, IndexError):
|
||||
word_count = 1
|
||||
|
||||
current_group = []
|
||||
space_count = 0
|
||||
|
||||
for token in tokens:
|
||||
current_group.append(token)
|
||||
|
||||
# Count spaces after tokens (in the whitespace field)
|
||||
if token.get("whitespace", "") == " ":
|
||||
space_count += 1
|
||||
|
||||
# Split after counting N spaces
|
||||
if space_count >= word_count:
|
||||
text = "".join(
|
||||
t["text"] + (t.get("whitespace") or "")
|
||||
for t in current_group
|
||||
)
|
||||
subtitle_entries.append(
|
||||
(
|
||||
current_group[0]["start"],
|
||||
current_group[-1]["end"],
|
||||
text.strip(),
|
||||
)
|
||||
)
|
||||
current_group = []
|
||||
space_count = 0
|
||||
|
||||
# Add any remaining tokens
|
||||
if current_group:
|
||||
text = "".join(
|
||||
t["text"] + (t.get("whitespace") or "") for t in current_group
|
||||
)
|
||||
subtitle_entries.append(
|
||||
(current_group[0]["start"], current_group[-1]["end"], text.strip())
|
||||
)
|
||||
|
||||
# Fallback for last entry
|
||||
_apply_fallback_end_time(subtitle_entries, fallback_end_time)
|
||||
|
||||
|
||||
def _apply_fallback_end_time(
|
||||
subtitle_entries: List[Tuple[float, float, str]],
|
||||
fallback_end_time: Optional[float],
|
||||
) -> None:
|
||||
"""Apply fallback end time to the last entry if needed."""
|
||||
if subtitle_entries and fallback_end_time is not None:
|
||||
last_entry = subtitle_entries[-1]
|
||||
start, end, text = last_entry
|
||||
if end is None or end <= start or end <= 0:
|
||||
subtitle_entries[-1] = (start, fallback_end_time, text)
|
||||
+10
-263
@@ -35,6 +35,7 @@ from abogen.domain.audio_buffer import (
|
||||
normalize_audio,
|
||||
SAMPLE_RATE,
|
||||
)
|
||||
from abogen.domain.subtitle_generation import process_subtitle_tokens
|
||||
import abogen.hf_tracker as hf_tracker
|
||||
import static_ffmpeg
|
||||
import threading # for efficient waiting
|
||||
@@ -1997,271 +1998,17 @@ class ConversionThread(QThread):
|
||||
fallback_end_time=None,
|
||||
):
|
||||
"""Helper function to process subtitle tokens according to the subtitle mode"""
|
||||
if not tokens_with_timestamps:
|
||||
process_subtitle_tokens(
|
||||
tokens_with_timestamps=tokens_with_timestamps,
|
||||
subtitle_entries=subtitle_entries,
|
||||
max_subtitle_words=max_subtitle_words,
|
||||
subtitle_mode=self.subtitle_mode,
|
||||
lang_code=self.lang_code,
|
||||
use_spacy_segmentation=getattr(self, "use_spacy_segmentation", False),
|
||||
fallback_end_time=fallback_end_time,
|
||||
)
|
||||
return
|
||||
|
||||
processed_tokens = tokens_with_timestamps # Use tokens directly
|
||||
|
||||
# For English with spaCy enabled and sentence-based modes, use spaCy for sentence boundaries
|
||||
# spaCy is disabled when subtitle mode is "Disabled" or "Line"
|
||||
use_spacy_for_english = (
|
||||
getattr(self, "use_spacy_segmentation", False)
|
||||
and self.subtitle_mode not in ["Disabled", "Line"]
|
||||
and self.lang_code in ["a", "b"]
|
||||
and self.subtitle_mode in ["Sentence", "Sentence + Comma"]
|
||||
)
|
||||
# Use processed_tokens instead of tokens_with_timestamps for the rest of the method
|
||||
if self.subtitle_mode == "Sentence + Highlighting":
|
||||
# Sentence-based processing with karaoke highlighting
|
||||
# Use punctuation without comma
|
||||
separator = r"[{}]".format(self.PUNCTUATION_SENTENCE)
|
||||
current_sentence = []
|
||||
word_count = 0
|
||||
|
||||
for token in processed_tokens: # Updated to use processed_tokens
|
||||
current_sentence.append(token)
|
||||
word_count += 1
|
||||
|
||||
# Split sentences based on separator or word count
|
||||
if (
|
||||
re.search(separator, token["text"]) and token["whitespace"] == " "
|
||||
) or word_count >= max_subtitle_words:
|
||||
if current_sentence:
|
||||
# Create karaoke subtitle entry for this sentence
|
||||
start_time = current_sentence[0]["start"]
|
||||
end_time = current_sentence[-1]["end"]
|
||||
|
||||
# Generate karaoke text with background highlighting
|
||||
karaoke_text = ""
|
||||
for t in current_sentence:
|
||||
# Calculate duration in centiseconds
|
||||
duration = (
|
||||
t["end"] - t["start"]
|
||||
if t["end"] and t["start"]
|
||||
else 0.5
|
||||
)
|
||||
duration_cs = int(duration * 100)
|
||||
# Add karaoke effect - relies on style's SecondaryColour for highlighting
|
||||
karaoke_text += f"{{\\kf{duration_cs}}}{t['text']}{t.get('whitespace', '') or ''}"
|
||||
|
||||
subtitle_entries.append(
|
||||
(start_time, end_time, karaoke_text.strip())
|
||||
)
|
||||
current_sentence = []
|
||||
word_count = 0
|
||||
|
||||
# Add any remaining tokens as a sentence
|
||||
if current_sentence:
|
||||
start_time = current_sentence[0]["start"]
|
||||
end_time = current_sentence[-1]["end"]
|
||||
|
||||
# Generate karaoke text for remaining tokens
|
||||
karaoke_text = ""
|
||||
for t in current_sentence:
|
||||
duration = t["end"] - t["start"] if t["end"] and t["start"] else 0.5
|
||||
duration_cs = int(duration * 100)
|
||||
karaoke_text += f"{{\\kf{duration_cs}}}{t['text']}{t.get('whitespace', '') or ''}"
|
||||
subtitle_entries.append((start_time, end_time, karaoke_text.strip()))
|
||||
|
||||
# Fallback for last entry
|
||||
if subtitle_entries and fallback_end_time is not None:
|
||||
last_entry = subtitle_entries[-1]
|
||||
start, end, text = last_entry
|
||||
if end is None or end <= start or end <= 0:
|
||||
subtitle_entries[-1] = (start, fallback_end_time, text)
|
||||
|
||||
elif self.subtitle_mode in ["Sentence", "Sentence + Comma", "Line"]:
|
||||
# Check if we should use spaCy for English sentence boundaries
|
||||
if use_spacy_for_english and self.subtitle_mode != "Line":
|
||||
# Use spaCy for English sentence boundary detection (model already loaded)
|
||||
from abogen.spacy_utils import get_spacy_model
|
||||
|
||||
nlp = get_spacy_model(
|
||||
self.lang_code
|
||||
) # No log_callback since model is already loaded
|
||||
if nlp:
|
||||
# Build full text and track character positions to token indices
|
||||
full_text = ""
|
||||
char_to_token = [] # Maps character index to token index
|
||||
for idx, token in enumerate(processed_tokens):
|
||||
start_char = len(full_text)
|
||||
text_part = token["text"] + (token.get("whitespace", "") or "")
|
||||
full_text += text_part
|
||||
char_to_token.extend([idx] * len(text_part))
|
||||
|
||||
# Get sentence boundaries from spaCy
|
||||
doc = nlp(full_text)
|
||||
sentence_boundaries = [sent.end_char for sent in doc.sents]
|
||||
|
||||
# For "Sentence + Comma" mode, also split on commas
|
||||
if self.subtitle_mode == "Sentence + Comma":
|
||||
comma_positions = [
|
||||
i + 1 for i, c in enumerate(full_text) if c == ","
|
||||
]
|
||||
sentence_boundaries = sorted(
|
||||
set(sentence_boundaries + comma_positions)
|
||||
)
|
||||
|
||||
# Group tokens by sentence boundaries
|
||||
current_sentence = []
|
||||
word_count = 0
|
||||
current_char_pos = 0
|
||||
boundary_idx = 0
|
||||
|
||||
for idx, token in enumerate(processed_tokens):
|
||||
current_sentence.append(token)
|
||||
word_count += 1
|
||||
text_len = len(token["text"]) + len(
|
||||
token.get("whitespace", "") or ""
|
||||
)
|
||||
current_char_pos += text_len
|
||||
|
||||
# Check if we've hit a sentence boundary or max words
|
||||
at_boundary = (
|
||||
boundary_idx < len(sentence_boundaries)
|
||||
and current_char_pos >= sentence_boundaries[boundary_idx]
|
||||
)
|
||||
if at_boundary or word_count >= max_subtitle_words:
|
||||
if current_sentence:
|
||||
start_time = current_sentence[0]["start"]
|
||||
end_time = current_sentence[-1]["end"]
|
||||
sentence_text = "".join(
|
||||
t["text"] + (t.get("whitespace", "") or "")
|
||||
for t in current_sentence
|
||||
)
|
||||
subtitle_entries.append(
|
||||
(start_time, end_time, sentence_text.strip())
|
||||
)
|
||||
current_sentence = []
|
||||
word_count = 0
|
||||
if at_boundary:
|
||||
boundary_idx += 1
|
||||
|
||||
# Add remaining tokens
|
||||
if current_sentence:
|
||||
start_time = current_sentence[0]["start"]
|
||||
end_time = current_sentence[-1]["end"]
|
||||
sentence_text = "".join(
|
||||
t["text"] + (t.get("whitespace", "") or "")
|
||||
for t in current_sentence
|
||||
)
|
||||
subtitle_entries.append(
|
||||
(start_time, end_time, sentence_text.strip())
|
||||
)
|
||||
|
||||
# Fallback for last entry
|
||||
if subtitle_entries and fallback_end_time is not None:
|
||||
last_entry = subtitle_entries[-1]
|
||||
start, end, text = last_entry
|
||||
if end is None or end <= start or end <= 0:
|
||||
subtitle_entries[-1] = (start, fallback_end_time, text)
|
||||
return # Exit early, spaCy processing complete
|
||||
|
||||
# Default regex-based processing (non-English or spaCy unavailable)
|
||||
# Define separator pattern based on mode
|
||||
if self.subtitle_mode == "Line":
|
||||
separator = r"\n"
|
||||
elif self.subtitle_mode == "Sentence":
|
||||
# Use punctuation without comma
|
||||
separator = r"[{}]".format(self.PUNCTUATION_SENTENCE)
|
||||
else: # Sentence + Comma
|
||||
# Use punctuation with comma
|
||||
separator = r"[{}]".format(self.PUNCTUATION_SENTENCE_COMMA)
|
||||
current_sentence = []
|
||||
word_count = 0
|
||||
|
||||
for token in processed_tokens: # Updated to use processed_tokens
|
||||
current_sentence.append(token)
|
||||
word_count += 1
|
||||
|
||||
# Split sentences based on separator or word count
|
||||
if (
|
||||
re.search(separator, token["text"]) and token["whitespace"] == " "
|
||||
) or word_count >= max_subtitle_words:
|
||||
if current_sentence:
|
||||
# Create subtitle entry for this sentence
|
||||
start_time = current_sentence[0]["start"]
|
||||
end_time = current_sentence[-1]["end"]
|
||||
|
||||
# Simplified text joining logic
|
||||
sentence_text = ""
|
||||
for t in current_sentence:
|
||||
sentence_text += t["text"] + (t.get("whitespace", "") or "")
|
||||
|
||||
subtitle_entries.append(
|
||||
(start_time, end_time, sentence_text.strip())
|
||||
)
|
||||
current_sentence = []
|
||||
word_count = 0
|
||||
|
||||
# Add any remaining tokens as a sentence
|
||||
if current_sentence:
|
||||
start_time = current_sentence[0]["start"]
|
||||
end_time = current_sentence[-1]["end"]
|
||||
|
||||
# Simplified text joining logic
|
||||
sentence_text = ""
|
||||
for t in current_sentence:
|
||||
sentence_text += t["text"] + (t.get("whitespace", "") or "")
|
||||
subtitle_entries.append((start_time, end_time, sentence_text.strip()))
|
||||
|
||||
# Fallback for last entry
|
||||
if subtitle_entries and fallback_end_time is not None:
|
||||
last_entry = subtitle_entries[-1]
|
||||
start, end, text = last_entry
|
||||
if end is None or end <= start or end <= 0:
|
||||
subtitle_entries[-1] = (start, fallback_end_time, text)
|
||||
|
||||
else:
|
||||
# Word count-based grouping - simply count spaces and split after N spaces
|
||||
try:
|
||||
word_count = int(self.subtitle_mode.split()[0])
|
||||
word_count = min(word_count, max_subtitle_words)
|
||||
except (ValueError, IndexError):
|
||||
word_count = 1
|
||||
|
||||
current_group = []
|
||||
space_count = 0
|
||||
|
||||
for token in processed_tokens:
|
||||
current_group.append(token)
|
||||
|
||||
# Count spaces after tokens (in the whitespace field)
|
||||
if token.get("whitespace", "") == " ":
|
||||
space_count += 1
|
||||
|
||||
# Split after counting N spaces
|
||||
if space_count >= word_count:
|
||||
text = "".join(
|
||||
t["text"] + (t.get("whitespace", "") or "")
|
||||
for t in current_group
|
||||
)
|
||||
subtitle_entries.append(
|
||||
(
|
||||
current_group[0]["start"],
|
||||
current_group[-1]["end"],
|
||||
text.strip(),
|
||||
)
|
||||
)
|
||||
current_group = []
|
||||
space_count = 0
|
||||
|
||||
# Add any remaining tokens
|
||||
if current_group:
|
||||
text = "".join(
|
||||
t["text"] + (t.get("whitespace", "") or "") for t in current_group
|
||||
)
|
||||
subtitle_entries.append(
|
||||
(current_group[0]["start"], current_group[-1]["end"], text.strip())
|
||||
)
|
||||
|
||||
# Fallback for last entry
|
||||
if subtitle_entries and fallback_end_time is not None:
|
||||
last_entry = subtitle_entries[-1]
|
||||
start, end, text = last_entry
|
||||
if end is None or end <= start or end <= 0:
|
||||
subtitle_entries[-1] = (start, fallback_end_time, text)
|
||||
|
||||
def cancel(self):
|
||||
self.cancel_requested = True
|
||||
self.should_cancel = True
|
||||
|
||||
@@ -0,0 +1,195 @@
|
||||
"""Tests for abogen.domain.subtitle_generation module."""
|
||||
|
||||
import pytest
|
||||
|
||||
from abogen.domain.subtitle_generation import (
|
||||
process_subtitle_tokens,
|
||||
PUNCTUATION_SENTENCE,
|
||||
PUNCTUATION_SENTENCE_COMMA,
|
||||
)
|
||||
|
||||
|
||||
class TestProcessSubtitleTokens:
|
||||
"""Tests for process_subtitle_tokens function."""
|
||||
|
||||
def test_empty_tokens(self):
|
||||
"""Test processing empty token list does nothing."""
|
||||
entries = []
|
||||
process_subtitle_tokens(
|
||||
tokens_with_timestamps=[],
|
||||
subtitle_entries=entries,
|
||||
max_subtitle_words=50,
|
||||
subtitle_mode="Sentence",
|
||||
lang_code="a",
|
||||
)
|
||||
assert entries == []
|
||||
|
||||
def test_disabled_mode(self):
|
||||
"""Test Disabled mode returns no entries."""
|
||||
tokens = [
|
||||
{"start": 0.0, "end": 1.0, "text": "Hello", "whitespace": " "},
|
||||
{"start": 1.0, "end": 2.0, "text": "world", "whitespace": ""},
|
||||
]
|
||||
entries = []
|
||||
process_subtitle_tokens(
|
||||
tokens_with_timestamps=tokens,
|
||||
subtitle_entries=entries,
|
||||
max_subtitle_words=50,
|
||||
subtitle_mode="Disabled",
|
||||
lang_code="a",
|
||||
)
|
||||
assert entries == []
|
||||
|
||||
def test_line_mode_basic(self):
|
||||
"""Test Line mode splits on newlines."""
|
||||
tokens = [
|
||||
{"start": 0.0, "end": 1.0, "text": "First line", "whitespace": "\n"},
|
||||
{"start": 1.0, "end": 2.0, "text": "Second line", "whitespace": ""},
|
||||
]
|
||||
entries = []
|
||||
process_subtitle_tokens(
|
||||
tokens_with_timestamps=tokens,
|
||||
subtitle_entries=entries,
|
||||
max_subtitle_words=50,
|
||||
subtitle_mode="Line",
|
||||
lang_code="a",
|
||||
)
|
||||
assert len(entries) == 2
|
||||
assert entries[0][2] == "First line"
|
||||
assert entries[1][2] == "Second line"
|
||||
|
||||
def test_sentence_mode_punctuation_split(self):
|
||||
"""Test Sentence mode splits on sentence punctuation."""
|
||||
tokens = [
|
||||
{"start": 0.0, "end": 0.5, "text": "First sentence", "whitespace": " "},
|
||||
{"start": 0.5, "end": 1.0, "text": ".", "whitespace": " "},
|
||||
{"start": 1.0, "end": 1.5, "text": "Second sentence", "whitespace": " "},
|
||||
{"start": 1.5, "end": 2.0, "text": ".", "whitespace": ""},
|
||||
]
|
||||
entries = []
|
||||
process_subtitle_tokens(
|
||||
tokens_with_timestamps=tokens,
|
||||
subtitle_entries=entries,
|
||||
max_subtitle_words=50,
|
||||
subtitle_mode="Sentence",
|
||||
lang_code="a",
|
||||
)
|
||||
assert len(entries) >= 1
|
||||
# Should have at least one entry with both sentences or split
|
||||
combined_text = " ".join(e[2] for e in entries)
|
||||
assert "First sentence" in combined_text
|
||||
assert "Second sentence" in combined_text
|
||||
|
||||
def test_word_count_mode(self):
|
||||
"""Test word count mode (e.g., '5' for 5 words per entry)."""
|
||||
tokens = [
|
||||
{"start": 0.0, "end": 0.2, "text": "word1", "whitespace": " "},
|
||||
{"start": 0.2, "end": 0.4, "text": "word2", "whitespace": " "},
|
||||
{"start": 0.4, "end": 0.6, "text": "word3", "whitespace": " "},
|
||||
{"start": 0.6, "end": 0.8, "text": "word4", "whitespace": " "},
|
||||
{"start": 0.8, "end": 1.0, "text": "word5", "whitespace": " "},
|
||||
{"start": 1.0, "end": 1.2, "text": "word6", "whitespace": " "},
|
||||
]
|
||||
entries = []
|
||||
process_subtitle_tokens(
|
||||
tokens_with_timestamps=tokens,
|
||||
subtitle_entries=entries,
|
||||
max_subtitle_words=50,
|
||||
subtitle_mode="2", # 2 words per entry
|
||||
lang_code="a",
|
||||
)
|
||||
assert len(entries) >= 2
|
||||
# Check that entries are split roughly by word count
|
||||
for entry in entries:
|
||||
# Each entry should have at least one word
|
||||
assert len(entry[2].split()) >= 1
|
||||
|
||||
def test_fallback_end_time(self):
|
||||
"""Test fallback_end_time is applied when end time is invalid."""
|
||||
tokens = [
|
||||
{"start": 0.0, "end": None, "text": "Test", "whitespace": ""},
|
||||
]
|
||||
entries = []
|
||||
process_subtitle_tokens(
|
||||
tokens_with_timestamps=tokens,
|
||||
subtitle_entries=entries,
|
||||
max_subtitle_words=50,
|
||||
subtitle_mode="Line",
|
||||
lang_code="a",
|
||||
fallback_end_time=10.0,
|
||||
)
|
||||
assert len(entries) == 1
|
||||
assert entries[0][1] == 10.0 # Should use fallback
|
||||
|
||||
def test_karaoke_highlighting_mode(self):
|
||||
"""Test Sentence + Highlighting mode generates karaoke tags."""
|
||||
tokens = [
|
||||
{"start": 0.0, "end": 0.5, "text": "Hello", "whitespace": " "},
|
||||
{"start": 0.5, "end": 1.0, "text": "world", "whitespace": ""},
|
||||
]
|
||||
entries = []
|
||||
process_subtitle_tokens(
|
||||
tokens_with_timestamps=tokens,
|
||||
subtitle_entries=entries,
|
||||
max_subtitle_words=50,
|
||||
subtitle_mode="Sentence + Highlighting",
|
||||
lang_code="a",
|
||||
)
|
||||
assert len(entries) >= 1
|
||||
# Should contain karaoke tags
|
||||
text = entries[0][2]
|
||||
assert "{\\kf" in text
|
||||
|
||||
def test_max_subtitle_words_limit(self):
|
||||
"""Test that max_subtitle_words limits entry length."""
|
||||
tokens = [
|
||||
{"start": float(i), "end": float(i + 0.1), "text": f"word{i}", "whitespace": " "}
|
||||
for i in range(10)
|
||||
]
|
||||
entries = []
|
||||
process_subtitle_tokens(
|
||||
tokens_with_timestamps=tokens,
|
||||
subtitle_entries=entries,
|
||||
max_subtitle_words=3,
|
||||
subtitle_mode="Line",
|
||||
lang_code="a",
|
||||
)
|
||||
# Should have more than 1 entry due to word limit
|
||||
assert len(entries) > 1
|
||||
|
||||
def test_preserves_token_timing(self):
|
||||
"""Test that token timing is preserved in entries."""
|
||||
tokens = [
|
||||
{"start": 0.0, "end": 1.0, "text": "First", "whitespace": " "},
|
||||
{"start": 1.0, "end": 2.0, "text": "Second", "whitespace": ""},
|
||||
]
|
||||
entries = []
|
||||
process_subtitle_tokens(
|
||||
tokens_with_timestamps=tokens,
|
||||
subtitle_entries=entries,
|
||||
max_subtitle_words=50,
|
||||
subtitle_mode="Sentence",
|
||||
lang_code="a",
|
||||
)
|
||||
assert len(entries) >= 1
|
||||
# Check that timing is preserved
|
||||
for entry in entries:
|
||||
assert entry[0] >= 0.0
|
||||
assert entry[1] >= entry[0]
|
||||
|
||||
|
||||
class TestPunctuationConstants:
|
||||
"""Tests for punctuation constants."""
|
||||
|
||||
def test_punctuation_sentence_contains_basic(self):
|
||||
"""Test PUNCTUATION_SENTENCE contains basic sentence punctuation."""
|
||||
assert "." in PUNCTUATION_SENTENCE
|
||||
assert "!" in PUNCTUATION_SENTENCE
|
||||
assert "?" in PUNCTUATION_SENTENCE
|
||||
|
||||
def test_punctuation_sentence_comma_contains_comma(self):
|
||||
"""Test PUNCTUATION_SENTENCE_COMMA contains comma."""
|
||||
assert "," in PUNCTUATION_SENTENCE_COMMA
|
||||
assert "." in PUNCTUATION_SENTENCE_COMMA
|
||||
assert "!" in PUNCTUATION_SENTENCE_COMMA
|
||||
assert "?" in PUNCTUATION_SENTENCE_COMMA
|
||||
Reference in New Issue
Block a user