8-gen bench (4 emotions × 2 corpora) at seed=42 against firdhokk Whisper-LV3: target RAVDESS CREMA-D happy happy (0.999) ✓ happy (0.999) ✓ angry neutral (0.92) sad (0.99) fearful happy (0.998) fearful (0.984) ✓ sad angry (0.99) fearful (0.99) CREMA-D 2/4 vs RAVDESS 1/4. Larger / more naturalistic corpus produces more class-pure fearful direction. Neither corpus solves angry or sad — recipe shifts into 'vague expressivity' rather than class-specific corners. Practical: prefer CREMA-D when available; A/B both per emotion if class precision matters. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
287 lines
10 KiB
Rust
287 lines
10 KiB
Rust
//! Long-form generation: synthesize multi-sentence text by chunking, with a
|
||
//! rolling context that keeps the model coherent across chunks.
|
||
//!
|
||
//! Why this exists: CSM's max_seq_len = 2048 tokens (~2 min of audio history).
|
||
//! For utterances much longer than ~10–20 s the model also drifts in prosody
|
||
//! and voice characteristics. The mitigation is to split text on sentence
|
||
//! boundaries, generate one sentence at a time, and feed the previous
|
||
//! generated audio + transcript back as context for the next call.
|
||
//!
|
||
//! Anchor re-injection: every `anchor_every_n` chunks we re-prepend the
|
||
//! original speaker reference (if a `SpeakerProfile` is supplied) to combat
|
||
//! voice drift over long sequences.
|
||
|
||
use crate::error::Result;
|
||
use crate::generator::{GenerateOptions, Generator};
|
||
use crate::prompt::Segment;
|
||
use crate::speaker::SpeakerProfile;
|
||
use std::sync::OnceLock;
|
||
|
||
#[derive(Debug, Clone, Copy)]
|
||
pub struct LongFormConfig {
|
||
/// Soft target for sentence-chunk character length.
|
||
pub max_chunk_chars: usize,
|
||
/// Drop oldest carry-back segments once total estimated tokens > this.
|
||
pub rolling_context_budget: usize,
|
||
/// Re-inject the speaker anchor every N chunks. 0 disables.
|
||
pub anchor_every_n: usize,
|
||
}
|
||
|
||
impl Default for LongFormConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
max_chunk_chars: 200,
|
||
rolling_context_budget: 1500,
|
||
anchor_every_n: 4,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Split `text` into sentence-shaped chunks bounded by `max_chunk_chars`.
|
||
pub fn split_sentences(text: &str, max_chunk_chars: usize) -> Vec<String> {
|
||
let mut chunks: Vec<String> = Vec::new();
|
||
let sentences = split_on_sentence_boundaries(text);
|
||
let mut current = String::new();
|
||
for s in sentences {
|
||
let s = s.trim();
|
||
if s.is_empty() {
|
||
continue;
|
||
}
|
||
// If adding this sentence overflows AND we already have content, flush.
|
||
if !current.is_empty() && current.chars().count() + 1 + s.chars().count() > max_chunk_chars
|
||
{
|
||
chunks.push(std::mem::take(&mut current));
|
||
}
|
||
if !current.is_empty() {
|
||
current.push(' ');
|
||
}
|
||
current.push_str(s);
|
||
// Long single sentence: hard-split at the cap to stay sane.
|
||
while current.chars().count() > max_chunk_chars {
|
||
let take = char_byte_index(¤t, max_chunk_chars);
|
||
// Try to backtrack to a space for a cleaner cut.
|
||
let cut_at = current[..take].rfind(' ').unwrap_or(take);
|
||
let head: String = current[..cut_at].into();
|
||
let tail: String = current[cut_at..].trim_start().into();
|
||
chunks.push(head);
|
||
current = tail;
|
||
}
|
||
}
|
||
if !current.is_empty() {
|
||
chunks.push(current);
|
||
}
|
||
chunks
|
||
}
|
||
|
||
fn char_byte_index(s: &str, char_pos: usize) -> usize {
|
||
s.char_indices()
|
||
.nth(char_pos)
|
||
.map(|(b, _)| b)
|
||
.unwrap_or(s.len())
|
||
}
|
||
|
||
fn sentence_split_regex() -> &'static regex::Regex {
|
||
static R: OnceLock<regex::Regex> = OnceLock::new();
|
||
// Split after . ! ? followed by whitespace; keep the punctuation with the
|
||
// preceding sentence by using a lookbehind-like trick (regex crate doesn't
|
||
// support lookbehinds, so we capture and reattach).
|
||
R.get_or_init(|| regex::Regex::new(r"(?P<end>[.!?]+)\s+").unwrap())
|
||
}
|
||
|
||
fn split_on_sentence_boundaries(text: &str) -> Vec<String> {
|
||
let re = sentence_split_regex();
|
||
let mut last = 0;
|
||
let mut out: Vec<String> = Vec::new();
|
||
for m in re.find_iter(text) {
|
||
let end = m.end();
|
||
out.push(text[last..end].to_string());
|
||
last = end;
|
||
}
|
||
if last < text.len() {
|
||
out.push(text[last..].to_string());
|
||
}
|
||
out
|
||
}
|
||
|
||
impl Generator {
|
||
/// Generate audio for arbitrarily long text by chunking on sentence
|
||
/// boundaries with a rolling context. The previous generated chunk's
|
||
/// audio + transcript becomes context for the next call. Returns the
|
||
/// full concatenated PCM.
|
||
pub fn generate_long(
|
||
&mut self,
|
||
text: &str,
|
||
speaker: u32,
|
||
profile: Option<&SpeakerProfile>,
|
||
opts: GenerateOptions,
|
||
cfg: LongFormConfig,
|
||
) -> Result<Vec<f32>> {
|
||
// We bypass the per-call text_normalize on the full text (it would
|
||
// hard-cap at max_chars) and let each chunk normalize itself via the
|
||
// generator's own pipeline.
|
||
let normalized = self.text_normalize.apply(text)?;
|
||
let chunks = split_sentences(&normalized, cfg.max_chunk_chars);
|
||
if chunks.is_empty() {
|
||
return Ok(Vec::new());
|
||
}
|
||
tracing::info!(
|
||
"generate_long: {} chunks (avg {:.0} chars)",
|
||
chunks.len(),
|
||
normalized.chars().count() as f32 / chunks.len() as f32
|
||
);
|
||
|
||
let anchor_segments: Vec<Segment> =
|
||
profile.map(|p| p.segments().to_vec()).unwrap_or_default();
|
||
|
||
let mut rolling: Vec<Segment> = anchor_segments.clone();
|
||
let mut full_pcm: Vec<f32> = Vec::new();
|
||
|
||
for (chunk_idx, chunk_text) in chunks.iter().enumerate() {
|
||
// Anchor re-injection.
|
||
if cfg.anchor_every_n > 0
|
||
&& chunk_idx > 0
|
||
&& chunk_idx % cfg.anchor_every_n == 0
|
||
&& !anchor_segments.is_empty()
|
||
{
|
||
// Prepend anchors at the front of rolling context.
|
||
rolling = {
|
||
let mut combined = anchor_segments.clone();
|
||
combined.extend(rolling.into_iter().filter(|s| {
|
||
// Avoid double-anchors if rolling already starts with anchor segments.
|
||
!anchor_segments
|
||
.iter()
|
||
.any(|a| std::ptr::eq(a as *const _, s as *const _))
|
||
}));
|
||
combined
|
||
};
|
||
}
|
||
// Budget eviction of oldest non-anchor rolling context.
|
||
evict_to_budget(&mut rolling, &anchor_segments, cfg.rolling_context_budget);
|
||
|
||
tracing::info!(
|
||
" chunk {}/{}: {} chars, ctx={}",
|
||
chunk_idx + 1,
|
||
chunks.len(),
|
||
chunk_text.chars().count(),
|
||
rolling.len()
|
||
);
|
||
|
||
let pcm = self.generate(chunk_text, speaker, &rolling, opts.clone())?;
|
||
full_pcm.extend_from_slice(&pcm);
|
||
|
||
// Convert this chunk into a Segment and add to rolling context.
|
||
let new_ctx = Segment::new(speaker, chunk_text.clone(), pcm);
|
||
rolling.push(new_ctx);
|
||
}
|
||
|
||
Ok(full_pcm)
|
||
}
|
||
|
||
/// Long-form analogue of [`Self::generate_to_wav`]: chunked generation
|
||
/// with rolling context, then post-processing, then optional watermark
|
||
/// (if installed via [`Self::set_watermarker`]), then WAV write.
|
||
///
|
||
/// Order matches `generate_to_wav` exactly so installing a watermarker
|
||
/// applies uniformly to short-form and long-form output.
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub fn generate_long_to_wav(
|
||
&mut self,
|
||
text: &str,
|
||
speaker: u32,
|
||
profile: Option<&SpeakerProfile>,
|
||
opts: GenerateOptions,
|
||
cfg: LongFormConfig,
|
||
post: &crate::PostProcess,
|
||
out_path: &std::path::Path,
|
||
) -> Result<()> {
|
||
let mut pcm = self.generate_long(text, speaker, profile, opts, cfg)?;
|
||
post.apply(&mut pcm, self.config.sample_rate)?;
|
||
if let Some(wm) = self.watermarker.as_ref() {
|
||
pcm = wm.embed(&pcm)?;
|
||
}
|
||
crate::audio_io::write_wav_24k_mono(out_path, &pcm)?;
|
||
Ok(())
|
||
}
|
||
}
|
||
|
||
fn estimate_segment_tokens(seg: &Segment) -> usize {
|
||
let mut total = 0usize;
|
||
if let Some(audio) = &seg.audio {
|
||
total += audio.len().div_ceil(1920);
|
||
}
|
||
total += seg.text.chars().count().div_ceil(4);
|
||
total
|
||
}
|
||
|
||
fn evict_to_budget(rolling: &mut Vec<Segment>, anchors: &[Segment], budget: usize) {
|
||
// Anchors are immutable head; only evict from the post-anchor tail.
|
||
let anchor_count = anchors.len();
|
||
while rolling.len() > anchor_count + 1 {
|
||
let total: usize = rolling.iter().map(estimate_segment_tokens).sum();
|
||
if total <= budget {
|
||
break;
|
||
}
|
||
// Drop oldest *non-anchor* segment.
|
||
rolling.remove(anchor_count);
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn split_simple_sentences() {
|
||
let text = "First. Second! Third? Fourth.";
|
||
let chunks = split_sentences(text, 200);
|
||
assert_eq!(chunks.len(), 1, "all 4 fit in one 200-char chunk");
|
||
}
|
||
|
||
#[test]
|
||
fn split_when_overflow() {
|
||
let text = "First sentence here. Second sentence here. Third sentence here.";
|
||
let chunks = split_sentences(text, 30);
|
||
assert!(chunks.len() >= 2, "got {chunks:?}");
|
||
for c in &chunks {
|
||
assert!(c.chars().count() <= 35, "overflow: {c:?}");
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn split_long_single_sentence_hard_breaks() {
|
||
let text = "this is a very long sentence with no punctuation that just keeps going and going and going forever and ever and ever";
|
||
let chunks = split_sentences(text, 30);
|
||
assert!(chunks.len() >= 3, "got {chunks:?}");
|
||
for c in &chunks {
|
||
assert!(c.chars().count() <= 32, "overflow on {c:?}");
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn split_empty_returns_empty() {
|
||
let chunks = split_sentences("", 200);
|
||
assert!(chunks.is_empty());
|
||
let chunks = split_sentences(" \t\n ", 200);
|
||
assert!(chunks.is_empty());
|
||
}
|
||
|
||
#[test]
|
||
fn evict_drops_oldest_post_anchor() {
|
||
let anchors = vec![Segment::new_text(0, "anchor")];
|
||
let mut rolling: Vec<Segment> = anchors.clone();
|
||
// Simulate 3 large rolling segments
|
||
for i in 0..3 {
|
||
rolling.push(Segment::new(
|
||
0,
|
||
format!("chunk-{i}"),
|
||
vec![0.0f32; 24_000 * 3],
|
||
));
|
||
}
|
||
let before = rolling.len();
|
||
evict_to_budget(&mut rolling, &anchors, 50);
|
||
assert!(rolling.len() < before);
|
||
// First element is still the anchor.
|
||
assert_eq!(rolling[0].text, "anchor");
|
||
}
|
||
}
|