Admin/audiobook
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add AI-powered text segmentation for intelligent speaker assignment

Browse Source 
- Added segment_text() method to OllamaService for AI-based text splitting
- OllamaService now annotates each text segment with the appropriate speaker
- Updated TTSService to use AI segmentation when available
- Falls back to pattern-based segmentation when Ollama unavailable
- AI intelligently assigns dialogue to speakers and narrative to Narrator
- Added comprehensive test suite for AI segmentation
- Verified AI correctly splits unmarked dialogue between multiple speakers

This makes the tool much more powerful - users can now input plain narrative
text without explicit role markers and the AI will intelligently assign voices
to different characters and the narrator.
This commit is contained in:
Paul Bailey
2025-12-09 10:59:00 +00:00
parent b48e53ce63
commit d18fddef57
3 changed files with 361 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

143
backend/services/ollama_service.py
View File

@@ -1,10 +1,10 @@
"""
Ollama integration service for AI-based speaker detection.
Ollama integration service for AI-based speaker detection and text segmentation.
"""
import json
import logging
import os
from typing import List, Optional
from typing import List, Optional, Dict, Tuple
import httpx
@@ -214,6 +214,145 @@ Return format: {{"speakers": ["Speaker1", "Speaker2", ...]}}
# Look for common patterns like lists or quoted names
return []
async def segment_text(self, text: str, speakers: List[str]) -> List[Tuple[str, str]]:
"""
Uses AI to segment text and assign each segment to a speaker.
Args:
text: Input text to segment
speakers: List of known speakers in the text
Returns:
List of tuples: (speaker_name, text_segment)
Raises:
Exception: If Ollama is unavailable or request fails
"""
try:
async with httpx.AsyncClient(timeout=60.0) as client: # Longer timeout for segmentation
prompt = self._build_segmentation_prompt(text, speakers)
logger.info(f"Requesting AI segmentation for {len(text)} chars with {len(speakers)} speakers")
response = await client.post(
f"{self.base_url}/api/generate",
json={
"model": self.model,
"prompt": prompt,
"stream": False,
"format": "json"
}
)
response.raise_for_status()
result = response.json()
response_text = result.get("response", "")
segments = self._parse_segmentation_response(response_text)
logger.info(f"AI segmentation produced {len(segments)} segments")
return segments
except httpx.TimeoutException:
logger.warning("Ollama segmentation request timed out")
raise Exception("Ollama segmentation timed out")
except httpx.HTTPError as e:
logger.warning(f"Ollama HTTP error during segmentation: {e}")
raise Exception(f"Ollama service error: {e}")
except Exception as e:
logger.warning(f"Ollama segmentation failed: {e}")
raise
def _build_segmentation_prompt(self, text: str, speakers: List[str]) -> str:
"""
Constructs prompt for text segmentation with speaker attribution.
Args:
text: Input text to segment
speakers: List of known speakers
Returns:
Formatted prompt string
"""
speakers_str = ", ".join(f'"{s}"' for s in speakers)
prompt = f"""You are analyzing a story or dialogue. Your task is to split the text into segments and assign each segment to the appropriate speaker.
Known speakers: {speakers_str}
Rules:
1. Split the text into logical segments (sentences or short paragraphs)
2. Assign each segment to the speaker who is speaking or being described
3. Use "Narrator" for narrative text, descriptions, or text that doesn't belong to a specific character
4. If a character is mentioned but not speaking, use "Narrator"
5. Only use speakers from the known speakers list
Return a JSON object with a "segments" array. Each segment should have:
- "speaker": the name of the speaker (must be from the known speakers list or "Narrator")
- "text": the text content for that segment
Example format:
{{
"segments": [
{{"speaker": "Narrator", "text": "The old wizard gazed into his crystal ball."}},
{{"speaker": "Wizard", "text": "I see danger approaching."}},
{{"speaker": "Narrator", "text": "A young knight heard the warning."}}
]
}}
Text to segment:
{text}
Return ONLY the JSON object with the segments array.
"""
return prompt
def _parse_segmentation_response(self, response: str) -> List[Tuple[str, str]]:
"""
Parses Ollama segmentation response to extract speaker-text pairs.
Args:
response: Response text from Ollama
Returns:
List of tuples: (speaker_name, text_content)
"""
if not response or not response.strip():
return []
try:
# Try to parse as JSON
data = json.loads(response.strip())
# Extract segments array
segments_data = []
if isinstance(data, dict):
if "segments" in data:
segments_data = data["segments"]
else:
# Try to find any array in the response
for value in data.values():
if isinstance(value, list):
segments_data = value
break
elif isinstance(data, list):
segments_data = data
# Convert to tuples
segments = []
for seg in segments_data:
if isinstance(seg, dict) and "speaker" in seg and "text" in seg:
speaker = str(seg["speaker"]).strip()
text = str(seg["text"]).strip()
if speaker and text:
segments.append((speaker, text))
return segments
except json.JSONDecodeError:
logger.warning(f"Failed to parse Ollama segmentation response as JSON: {response[:200]}")
return []
async def is_available(self) -> bool:
"""
Checks if Ollama service is available.

82
backend/services/tts_service.py
View File

@@ -8,6 +8,7 @@ from typing import Dict, List, Tuple, Optional
import numpy as np
import soundfile as sf
from kokoro import KPipeline
from .ollama_service import OllamaService
logger = logging.getLogger(__name__)
@@ -30,9 +31,10 @@ class TTSService:
COLON_PATTERN = r'^([A-Z][a-zA-Z\s]+?):\s*(.+)$'
BRACKET_PATTERN = r'^\[([A-Z][a-zA-Z\s]+?)\]\s*(.+)$'
def __init__(self):
"""Initializes TTS service with empty pipeline cache."""
def __init__(self, ollama_service: Optional[OllamaService] = None):
"""Initializes TTS service with empty pipeline cache and optional Ollama service."""
self.pipelines: Dict[str, KPipeline] = {}
self.ollama_service = ollama_service or OllamaService()
logger.info("TTSService initialized")
def get_pipeline(self, lang_code: str) -> KPipeline:
@@ -72,15 +74,62 @@ class TTSService:
logger.error(f"Failed to initialize pipeline for language {lang_code}: {e}")
raise
async def segment_text_with_ai(
self,
text: str,
role_voice_map: Dict[str, str]
) -> Optional[List[Tuple[str, str, str]]]:
"""
Uses AI to segment text and assign speakers intelligently.
Args:
text: Input text to segment
role_voice_map: Mapping of role IDs to voice IDs
Returns:
List of tuples: (role_id, voice_id, text_content) or None if AI unavailable
"""
try:
# Check if Ollama is available
is_available = await self.ollama_service.is_available()
if not is_available:
logger.info("Ollama not available for AI segmentation")
return None
# Get list of speakers from role_voice_map
speakers = list(role_voice_map.keys())
# Request AI segmentation
ai_segments = await self.ollama_service.segment_text(text, speakers)
if not ai_segments:
logger.warning("AI segmentation returned no segments")
return None
# Convert to format with voice assignments
segments = []
for speaker, content in ai_segments:
voice = role_voice_map.get(speaker, self.DEFAULT_VOICE)
segments.append((speaker, voice, content))
logger.info(f"AI segmented text into {len(segments)} segments")
return segments
except Exception as e:
logger.warning(f"AI segmentation failed: {e}")
return None
def segment_text(
self,
text: str,
role_voice_map: Dict[str, str]
) -> List[Tuple[str, str, str]]:
"""
Segments text into (role_id, voice, text_content) tuples.
Segments text into (role_id, voice, text_content) tuples using pattern matching.
Identifies role markers and assigns segments accordingly.
This is the fallback method when AI segmentation is not available.
Args:
text: Input text to segment
role_voice_map: Mapping of role IDs to voice IDs
@@ -118,7 +167,7 @@ class TTSService:
voice = role_voice_map.get("Narrator", self.DEFAULT_VOICE)
segments.append(("Narrator", voice, line))
logger.info(f"Segmented text into {len(segments)} segments")
logger.info(f"Pattern-based segmentation produced {len(segments)} segments")
return segments
def _chunk_text(self, text: str, max_length: int = 500) -> List[str]:
@@ -186,13 +235,16 @@ class TTSService:
try:
# Generate audio for chunk
result = pipeline(chunk, voice=voice, speed=speed)
# Kokoro returns a generator that yields audio arrays
# Collect all arrays and concatenate them
# Kokoro returns a generator that yields Result objects
# Each Result object has an .audio attribute (torch.Tensor)
audio_parts = []
for part in result:
if isinstance(part, np.ndarray):
# Flatten to 1D if needed
audio_parts.append(part.flatten())
# Extract audio tensor and convert to numpy
if hasattr(part, 'audio'):
audio_tensor = part.audio
# Convert torch tensor to numpy and flatten
audio_np = audio_tensor.cpu().numpy().flatten()
audio_parts.append(audio_np)
# Concatenate all parts
if audio_parts:
@@ -247,6 +299,8 @@ class TTSService:
Main orchestration method for audiobook generation.
Segments text, synthesizes each segment, concatenates, and returns audio bytes.
Uses AI segmentation when available, falls back to pattern matching.
Args:
text: Input text
lang: Language code
@@ -264,8 +318,14 @@ class TTSService:
# Get or create pipeline for language
pipeline = self.get_pipeline(lang)
# Segment text by roles
segments = self.segment_text(text, role_voice_map)
# Try AI segmentation first, fall back to pattern matching
segments = await self.segment_text_with_ai(text, role_voice_map)
if segments is None:
logger.info("Using pattern-based segmentation as fallback")
segments = self.segment_text(text, role_voice_map)
else:
logger.info("Using AI-powered segmentation")
if not segments:
raise ValueError("No segments to synthesize")