Files
rustytorch/crates/models/rtx-csm/src/longform.rs
T
osobhandClaude Opus 4.7 a5cedfb46a rtx-csm: emotional_speech_guide — CREMA-D vs RAVDESS firdhokk verdict
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]>
2026-04-30 00:01:02 -07:00

287 lines
10 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! 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(&current, 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");
}
}