Spaces:
Sleeping
Sleeping
| """Text processing utilities for the TTS API.""" | |
| import re | |
| from typing import List | |
| class TextChunker: | |
| """Server-side text chunking for optimal GPU processing.""" | |
| def __init__(self, max_chunk_size: int = 800, overlap_sentences: int = 0): | |
| """ | |
| Initialize the text chunker. | |
| Args: | |
| max_chunk_size: Maximum number of characters per chunk | |
| overlap_sentences: Number of sentences to overlap between chunks for continuity | |
| """ | |
| self.max_chunk_size = max_chunk_size | |
| self.overlap_sentences = overlap_sentences | |
| def chunk_text(self, text: str) -> List[str]: | |
| """ | |
| Break text into smaller chunks based on paragraphs and sentence boundaries. | |
| Args: | |
| text: The input text to chunk | |
| Returns: | |
| List of text chunks | |
| """ | |
| if not text or not text.strip(): | |
| return [] | |
| # Clean the text | |
| text = text.strip() | |
| # If text is within the limit, return as single chunk | |
| if len(text) <= self.max_chunk_size: | |
| return [text] | |
| chunks = [] | |
| # First, try to split by paragraphs | |
| paragraphs = self._split_into_paragraphs(text) | |
| current_chunk = "" | |
| for paragraph in paragraphs: | |
| # If adding this paragraph would exceed the limit | |
| if len(current_chunk) + len(paragraph) + 1 > self.max_chunk_size: | |
| # If we have content in current chunk, save it | |
| if current_chunk.strip(): | |
| chunks.append(current_chunk.strip()) | |
| current_chunk = "" | |
| # If the paragraph itself is too long, split it by sentences | |
| if len(paragraph) > self.max_chunk_size: | |
| sentence_chunks = self._split_paragraph_into_sentences(paragraph) | |
| for sentence_chunk in sentence_chunks: | |
| if len(current_chunk) + len(sentence_chunk) + 1 > self.max_chunk_size: | |
| if current_chunk.strip(): | |
| chunks.append(current_chunk.strip()) | |
| current_chunk = sentence_chunk | |
| else: | |
| if current_chunk: | |
| current_chunk += " " + sentence_chunk | |
| else: | |
| current_chunk = sentence_chunk | |
| else: | |
| current_chunk = paragraph | |
| else: | |
| # Add paragraph to current chunk | |
| if current_chunk: | |
| current_chunk += "\n\n" + paragraph | |
| else: | |
| current_chunk = paragraph | |
| # Add any remaining content | |
| if current_chunk.strip(): | |
| chunks.append(current_chunk.strip()) | |
| # Apply overlap if specified | |
| if self.overlap_sentences > 0 and len(chunks) > 1: | |
| chunks = self._add_overlap(chunks) | |
| return chunks | |
| def _split_into_paragraphs(self, text: str) -> List[str]: | |
| """Split text into paragraphs.""" | |
| # Split by double newlines or multiple spaces | |
| paragraphs = re.split(r'\n\s*\n|(?:\n\s*){2,}', text) | |
| # Filter out empty paragraphs and strip whitespace | |
| return [p.strip() for p in paragraphs if p.strip()] | |
| def _split_paragraph_into_sentences(self, paragraph: str) -> List[str]: | |
| """Split a long paragraph into sentence-based chunks.""" | |
| # Split by sentence boundaries | |
| sentences = re.split(r'(?<=[.!?])\s+', paragraph) | |
| chunks = [] | |
| current_chunk = "" | |
| for sentence in sentences: | |
| # If a single sentence is longer than max_chunk_size, we need to force-split it | |
| if len(sentence) > self.max_chunk_size: | |
| # Save current chunk if it has content | |
| if current_chunk.strip(): | |
| chunks.append(current_chunk.strip()) | |
| current_chunk = "" | |
| # Force-split the long sentence into smaller pieces | |
| while len(sentence) > self.max_chunk_size: | |
| # Find a good breaking point (prefer spaces) | |
| break_point = self.max_chunk_size | |
| if ' ' in sentence[:self.max_chunk_size]: | |
| # Find the last space within the limit | |
| break_point = sentence[:self.max_chunk_size].rfind(' ') | |
| chunk_part = sentence[:break_point] | |
| chunks.append(chunk_part) | |
| sentence = sentence[break_point:].strip() | |
| # Add the remaining part of the sentence | |
| if sentence: | |
| current_chunk = sentence | |
| elif len(current_chunk) + len(sentence) + 1 > self.max_chunk_size: | |
| if current_chunk.strip(): | |
| chunks.append(current_chunk.strip()) | |
| current_chunk = sentence | |
| else: | |
| if current_chunk: | |
| current_chunk += " " + sentence | |
| else: | |
| current_chunk = sentence | |
| if current_chunk.strip(): | |
| chunks.append(current_chunk.strip()) | |
| return chunks | |
| def _add_overlap(self, chunks: List[str]) -> List[str]: | |
| """Add sentence overlap between chunks for better continuity.""" | |
| if len(chunks) <= 1: | |
| return chunks | |
| overlapped_chunks = [chunks[0]] # First chunk stays the same | |
| for i in range(1, len(chunks)): | |
| # Get last few sentences from previous chunk | |
| prev_chunk = chunks[i - 1] | |
| current_chunk = chunks[i] | |
| prev_sentences = re.split(r'(?<=[.!?])\s+', prev_chunk) | |
| overlap_text = " ".join(prev_sentences[-self.overlap_sentences:]) if len(prev_sentences) > self.overlap_sentences else "" | |
| if overlap_text: | |
| overlapped_chunk = overlap_text + " " + current_chunk | |
| else: | |
| overlapped_chunk = current_chunk | |
| overlapped_chunks.append(overlapped_chunk) | |
| return overlapped_chunks | |
| def get_chunk_info(self, chunks: List[str]) -> dict: | |
| """Get information about the chunks.""" | |
| return { | |
| "total_chunks": len(chunks), | |
| "total_characters": sum(len(chunk) for chunk in chunks), | |
| "avg_chunk_size": sum(len(chunk) for chunk in chunks) / len(chunks) if chunks else 0, | |
| "max_chunk_size": max(len(chunk) for chunk in chunks) if chunks else 0, | |
| "min_chunk_size": min(len(chunk) for chunk in chunks) if chunks else 0 | |
| } | |