diodiogod-TTS-Audio-Suite/utils/text/chunking.py

"""
Text Chunking - Enhanced text chunker for ChatterBox Voice
Extracted from the original nodes.py implementation
"""

import re
import torch
from typing import List


class ImprovedChatterBoxChunker:
    """Enhanced text chunker inspired by Orpheus TTS approach"""
    
    @staticmethod
    def split_into_chunks(text: str, max_chars: int = 400) -> List[str]:
        """
        Split text into chunks with better sentence boundary handling.
        Uses character-based limits like Orpheus TTS for more predictable chunk sizes.
        
        Args:
            text: Input text to split
            max_chars: Maximum characters per chunk
            
        Returns:
            List of text chunks
        """
        if not text.strip():
            return []
            
        # Clean and normalize text
        text = re.sub(r'\s+', ' ', text.strip())
        
        # If text is short enough, return as single chunk
        if len(text) <= max_chars:
            return [text]
        
        # Split into sentences using robust regex (same as Orpheus)
        sentences = re.split(r'(?<=[.!?])\s+', text)
        
        chunks = []
        current_chunk = ""
        
        for sentence in sentences:
            sentence = sentence.strip()
            if not sentence:
                continue
            
            # If adding this sentence exceeds limit and we have content, start new chunk
            if len(current_chunk) + len(sentence) + 1 > max_chars and current_chunk:
                chunks.append(current_chunk.strip())
                current_chunk = sentence
            # If single sentence is too long, split it further
            elif len(sentence) > max_chars:
                # Add current chunk if not empty
                if current_chunk:
                    chunks.append(current_chunk.strip())
                    current_chunk = ""
                
                # Split long sentence by commas (Orpheus approach)
                parts = re.split(r'(?<=,)\s+', sentence)
                sub_chunk = ""
                
                for part in parts:
                    if len(sub_chunk) + len(part) + 1 > max_chars:
                        if sub_chunk:
                            chunks.append(sub_chunk.strip())
                            sub_chunk = part
                        else:
                            # Even single part is too long - split arbitrarily
                            for i in range(0, len(part), max_chars):
                                chunk_part = part[i:i+max_chars].strip()
                                if chunk_part:
                                    chunks.append(chunk_part)
                    else:
                        sub_chunk = sub_chunk + ", " + part if sub_chunk else part
                
                # Set remaining as current chunk
                if sub_chunk:
                    current_chunk = sub_chunk
            else:
                # Normal sentence - add to current chunk
                current_chunk = current_chunk + " " + sentence if current_chunk else sentence
        
        # Add final chunk
        if current_chunk.strip():
            chunks.append(current_chunk.strip())
        
        return chunks
    
    @staticmethod
    def add_silence_padding(audio: torch.Tensor, duration_ms: int = 50, sample_rate: int = 22050) -> torch.Tensor:
        """
        Add brief silence between chunks to improve naturalness
        
        Args:
            audio: Input audio tensor
            duration_ms: Duration of silence in milliseconds
            sample_rate: Sample rate of the audio
            
        Returns:
            Audio tensor with silence padding added
        """
        silence_samples = int(duration_ms * sample_rate / 1000)
        
        # Create silence tensor with same shape as audio tensor
        if audio.dim() == 1:
            # 1D audio tensor [samples]
            silence = torch.zeros(silence_samples, device=audio.device, dtype=audio.dtype)
        elif audio.dim() == 2:
            # 2D audio tensor [channels, samples]
            silence = torch.zeros(audio.shape[0], silence_samples, device=audio.device, dtype=audio.dtype)
        else:
            # Fallback - just match the last dimension
            silence_shape = list(audio.shape)
            silence_shape[-1] = silence_samples
            silence = torch.zeros(*silence_shape, device=audio.device, dtype=audio.dtype)
        
        return torch.cat([audio, silence], dim=-1)
    
    @staticmethod
    def estimate_chunk_count(text: str, max_chars: int = 400) -> int:
        """
        Estimate how many chunks the text will be split into without actually splitting.
        Useful for progress tracking.
        
        Args:
            text: Input text to estimate
            max_chars: Maximum characters per chunk
            
        Returns:
            Estimated number of chunks
        """
        if not text.strip():
            return 0
        
        text = re.sub(r'\s+', ' ', text.strip())
        
        if len(text) <= max_chars:
            return 1
        
        # Quick estimation based on average chunk size
        # This is approximate but good enough for progress estimation
        return max(1, len(text) // max_chars + (1 if len(text) % max_chars > 0 else 0))
    
    @staticmethod
    def get_chunk_stats(chunks: List[str]) -> dict:
        """
        Get statistics about the chunks.
        
        Args:
            chunks: List of text chunks
            
        Returns:
            Dictionary with chunk statistics
        """
        if not chunks:
            return {
                "total_chunks": 0,
                "total_chars": 0,
                "avg_chars_per_chunk": 0,
                "min_chars": 0,
                "max_chars": 0,
                "shortest_chunk": "",
                "longest_chunk": ""
            }
        
        chunk_lengths = [len(chunk) for chunk in chunks]
        
        return {
            "total_chunks": len(chunks),
            "total_chars": sum(chunk_lengths),
            "avg_chars_per_chunk": sum(chunk_lengths) / len(chunks),
            "min_chars": min(chunk_lengths),
            "max_chars": max(chunk_lengths),
            "shortest_chunk": min(chunks, key=len),
            "longest_chunk": max(chunks, key=len)
        }
    
    @classmethod
    def validate_chunking_params(cls, max_chars: int) -> int:
        """
        Validate and normalize chunking parameters.
        
        Args:
            max_chars: Maximum characters per chunk
            
        Returns:
            Validated max_chars value
        """
        # Ensure reasonable bounds
        if max_chars < 100:
            max_chars = 100
        elif max_chars > 2000:
            max_chars = 2000
        
        return max_chars
    
    @classmethod
    def chunk_with_overlap(cls, text: str, max_chars: int = 400, overlap_chars: int = 50) -> List[str]:
        """
        Split text into overlapping chunks for better context preservation.
        
        Args:
            text: Input text to split
            max_chars: Maximum characters per chunk
            overlap_chars: Number of characters to overlap between chunks
            
        Returns:
            List of overlapping text chunks
        """
        if not text.strip():
            return []
        
        # Clean and normalize text
        text = re.sub(r'\s+', ' ', text.strip())
        
        if len(text) <= max_chars:
            return [text]
        
        chunks = []
        start = 0
        
        while start < len(text):
            end = start + max_chars
            
            if end >= len(text):
                # Last chunk
                chunks.append(text[start:].strip())
                break
            
            # Try to find a good break point (sentence boundary)
            chunk_text = text[start:end]
            
            # Look for sentence endings near the end of the chunk
            sentence_breaks = [m.end() for m in re.finditer(r'[.!?]\s+', chunk_text)]
            if sentence_breaks:
                # Use the last sentence break as the end point
                actual_end = start + sentence_breaks[-1]
                chunks.append(text[start:actual_end].strip())
                start = actual_end - overlap_chars
            else:
                # No good break point found, split at max_chars
                chunks.append(chunk_text.strip())
                start = end - overlap_chars
            
            # Ensure we don't go backwards
            if start < 0:
                start = 0
        
        return [chunk for chunk in chunks if chunk.strip()]