cf0a8b71fd
- Moved SentenceTextChunker and SentenceTextChunkerConfig to the utils module for better organization. - Updated pytest.ini to include the current directory in the Python path. - Added a new utils module with shared utility helpers. - Adjusted import paths in the test files to reflect the new location of text chunking classes.
136 lines
4.3 KiB
Python
136 lines
4.3 KiB
Python
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass
|
||
|
||
|
||
SENTENCE_ENDING_PUNCTUATION = frozenset(".!?;。！？；")
|
||
SOFT_BREAK_PUNCTUATION = frozenset(",，、:：")
|
||
CLOSING_PUNCTUATION = frozenset("\"'”’)]}）】》」』")
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class SentenceTextChunkerConfig:
|
||
"""Configuration for streaming text chunks sent to TTS."""
|
||
|
||
min_chars: int = 1
|
||
max_chars: int = 80
|
||
use_soft_breaks: bool = True
|
||
|
||
|
||
class SentenceTextChunker:
|
||
"""Lightweight Pipecat-style sentence chunker for streaming TTS text.
|
||
|
||
The chunker waits for one non-whitespace lookahead character after sentence
|
||
punctuation before emitting a sentence. This avoids splitting too early when
|
||
a punctuation mark arrives at the end of a stream token.
|
||
"""
|
||
|
||
def __init__(self, config: SentenceTextChunkerConfig | None = None) -> None:
|
||
self._config = config or SentenceTextChunkerConfig()
|
||
self._buffer = ""
|
||
self._needs_lookahead = False
|
||
self._last_soft_break = -1
|
||
|
||
@property
|
||
def text(self) -> str:
|
||
return self._buffer.strip(" ")
|
||
|
||
def feed(self, text: str) -> list[str]:
|
||
"""Append streaming text and return chunks ready for TTS."""
|
||
chunks: list[str] = []
|
||
if not text:
|
||
return chunks
|
||
|
||
for char in text:
|
||
self._buffer += char
|
||
index = len(self._buffer) - 1
|
||
|
||
if char in SOFT_BREAK_PUNCTUATION:
|
||
self._last_soft_break = index + 1
|
||
|
||
if self._needs_lookahead and char.strip():
|
||
self._needs_lookahead = False
|
||
chunk = self._pop_sentence_chunk()
|
||
if chunk:
|
||
chunks.append(chunk)
|
||
continue
|
||
|
||
if char in SENTENCE_ENDING_PUNCTUATION and not self._is_decimal_point(index):
|
||
self._needs_lookahead = True
|
||
|
||
chunk = self._pop_soft_chunk_if_needed()
|
||
if chunk:
|
||
chunks.append(chunk)
|
||
|
||
return chunks
|
||
|
||
def flush(self) -> str | None:
|
||
"""Return any remaining buffered text at end of stream."""
|
||
if not self._buffer:
|
||
return None
|
||
chunk = self._buffer.strip(" ")
|
||
self.reset()
|
||
return chunk or None
|
||
|
||
def reset(self) -> None:
|
||
self._buffer = ""
|
||
self._needs_lookahead = False
|
||
self._last_soft_break = -1
|
||
|
||
def _pop_sentence_chunk(self) -> str | None:
|
||
end = self._sentence_end_index()
|
||
if end is None:
|
||
return None
|
||
chunk = self._buffer[:end].strip(" ")
|
||
if len(chunk) < self._config.min_chars:
|
||
return None
|
||
self._buffer = self._buffer[end:]
|
||
self._last_soft_break = self._find_last_soft_break()
|
||
return chunk
|
||
|
||
def _sentence_end_index(self) -> int | None:
|
||
index = 0
|
||
while index < len(self._buffer):
|
||
char = self._buffer[index]
|
||
if char in SENTENCE_ENDING_PUNCTUATION and not self._is_decimal_point(index):
|
||
end = index + 1
|
||
while end < len(self._buffer) and self._buffer[end] in CLOSING_PUNCTUATION:
|
||
end += 1
|
||
return end
|
||
index += 1
|
||
return None
|
||
|
||
def _pop_soft_chunk_if_needed(self) -> str | None:
|
||
if (
|
||
not self._config.use_soft_breaks
|
||
or self._config.max_chars <= 0
|
||
or len(self._buffer) < self._config.max_chars
|
||
or self._last_soft_break <= 0
|
||
):
|
||
return None
|
||
|
||
chunk = self._buffer[: self._last_soft_break].strip(" ")
|
||
if len(chunk) < self._config.min_chars:
|
||
return None
|
||
|
||
self._buffer = self._buffer[self._last_soft_break :]
|
||
self._last_soft_break = self._find_last_soft_break()
|
||
self._needs_lookahead = False
|
||
return chunk
|
||
|
||
def _find_last_soft_break(self) -> int:
|
||
for index in range(len(self._buffer) - 1, -1, -1):
|
||
if self._buffer[index] in SOFT_BREAK_PUNCTUATION:
|
||
return index + 1
|
||
return -1
|
||
|
||
def _is_decimal_point(self, index: int) -> bool:
|
||
if self._buffer[index] != ".":
|
||
return False
|
||
return (
|
||
index > 0
|
||
and index + 1 < len(self._buffer)
|
||
and self._buffer[index - 1].isdigit()
|
||
and self._buffer[index + 1].isdigit()
|
||
)
|