308 lines
11 KiB
Rust
308 lines
11 KiB
Rust
//! Sprint 2 Phase B — derive ActAdd-style steering vectors from an
|
||
//! emotion-labeled manifest by running teacher-forced forward passes
|
||
//! through the backbone and computing per-layer difference-of-means
|
||
//! between a target-emotion pool and a neutral-baseline pool.
|
||
//!
|
||
//! Inspired by EmoSteer-TTS (arXiv 2508.03543); see `src/steering.rs`
|
||
//! for the architectural deviation note (CSM is AR-token, the paper is
|
||
//! flow-matching DiT).
|
||
//!
|
||
//! Manifest format (JSONL): rows with `wav`, `transcript`, `emotion_tag`
|
||
//! (e.g. `[whisper]`, `[neutral]`, `[excited]`). Brackets are optional;
|
||
//! they're stripped for matching. `audio_to_manifest --auto-emotion-tag
|
||
//! --use-emotion2vec` produces this layout.
|
||
//!
|
||
//! Output: a safetensors file with keys `layer_<i>_steering`, shape
|
||
//! `(embed_dim,)`. Loadable by `examples/generate --steering-vec ...`.
|
||
//!
|
||
//! Usage:
|
||
//! ```bash
|
||
//! target/release/examples/steering_extract \
|
||
//! --in /tmp/voice_corpus/manifest_xxx/manifest.jsonl \
|
||
//! --target whisper \
|
||
//! --baseline neutral \
|
||
//! --out /tmp/whisper_steering.safetensors
|
||
//! ```
|
||
|
||
use anyhow::{Context, Result};
|
||
use candle_core::{DType, Device, Tensor};
|
||
use clap::Parser;
|
||
use rtx_csm::Generator;
|
||
use rtx_csm::audio_io;
|
||
use rtx_csm::prompt::{Segment, build_prompt};
|
||
use serde::Deserialize;
|
||
use std::collections::HashMap;
|
||
use std::path::PathBuf;
|
||
|
||
#[derive(Debug, Parser)]
|
||
struct Cli {
|
||
/// Manifest JSONL with `wav`, `transcript`, `emotion_tag` rows. The
|
||
/// manifest's wav paths can be absolute or relative to the manifest dir.
|
||
#[arg(long = "in")]
|
||
input: PathBuf,
|
||
/// Target emotion (e.g. "whisper", "excited", "sad"). Brackets are
|
||
/// optional — `whisper` and `[whisper]` both match.
|
||
#[arg(long)]
|
||
target: String,
|
||
/// Baseline emotion to subtract (default: "neutral").
|
||
#[arg(long, default_value = "neutral")]
|
||
baseline: String,
|
||
/// Output safetensors path.
|
||
#[arg(long)]
|
||
out: PathBuf,
|
||
/// Maximum rows per pool to encode. 0 = no cap. Lower for quick smoke.
|
||
#[arg(long, default_value_t = 50)]
|
||
max_per_pool: usize,
|
||
/// Which module's activations to capture: `backbone` (default; semantic
|
||
/// codebook 0) or `decoder` (acoustic codebooks 1..N-1; smaller embed
|
||
/// dim, designed to bias prosody/timbre without disturbing word
|
||
/// content). The decoder path teacher-forces the ground-truth c0 and
|
||
/// captures one frame per clip.
|
||
#[arg(long, default_value = "backbone")]
|
||
target_module: String,
|
||
/// Force CPU.
|
||
#[arg(long)]
|
||
cpu: bool,
|
||
}
|
||
|
||
#[derive(Debug, Deserialize)]
|
||
struct Row {
|
||
wav: PathBuf,
|
||
transcript: String,
|
||
#[serde(default)]
|
||
emotion_tag: Option<String>,
|
||
#[serde(default)]
|
||
speaker: Option<u32>,
|
||
}
|
||
|
||
fn normalize_tag(s: &str) -> String {
|
||
s.trim()
|
||
.trim_start_matches('[')
|
||
.trim_end_matches(']')
|
||
.to_lowercase()
|
||
}
|
||
|
||
fn main() -> Result<()> {
|
||
tracing_subscriber::fmt().with_target(false).init();
|
||
let cli = Cli::parse();
|
||
|
||
let device = if cli.cpu {
|
||
Device::Cpu
|
||
} else {
|
||
Generator::default_device()?
|
||
};
|
||
tracing::info!("device: {device:?}");
|
||
|
||
let mut generator = Generator::load_csm_1b(&device)?;
|
||
let module = cli.target_module.to_lowercase();
|
||
if !["backbone", "decoder"].contains(&module.as_str()) {
|
||
anyhow::bail!("--target-module must be 'backbone' or 'decoder' (got '{module}')");
|
||
}
|
||
let (num_layers, embed_dim): (usize, usize) = if module == "decoder" {
|
||
let n = generator
|
||
.decoder_num_layers()
|
||
.ok_or_else(|| anyhow::anyhow!("decoder steering needs the FP backbone"))?;
|
||
(n, 1024) // CSM-1B decoder = Llama100M
|
||
} else {
|
||
let n = generator
|
||
.backbone_num_layers()
|
||
.ok_or_else(|| anyhow::anyhow!("steering extraction requires the FP backbone"))?;
|
||
(n, 2048) // CSM-1B backbone = Llama1B
|
||
};
|
||
tracing::info!("loaded CSM-1B; module={module} → {num_layers} layers × {embed_dim} embed");
|
||
|
||
// Resolve manifest paths and normalize tags.
|
||
let manifest_dir = cli
|
||
.input
|
||
.parent()
|
||
.map(|p| p.to_path_buf())
|
||
.unwrap_or_else(|| PathBuf::from("."));
|
||
let target_key = normalize_tag(&cli.target);
|
||
let baseline_key = normalize_tag(&cli.baseline);
|
||
tracing::info!("target='{target_key}' baseline='{baseline_key}'");
|
||
|
||
let body = std::fs::read_to_string(&cli.input)
|
||
.with_context(|| format!("read manifest {}", cli.input.display()))?;
|
||
|
||
let mut pools: HashMap<String, (Vec<Vec<f32>>, usize)> = HashMap::new();
|
||
pools.insert(
|
||
target_key.clone(),
|
||
(vec![vec![0.0f32; embed_dim]; num_layers], 0),
|
||
);
|
||
pools.insert(
|
||
baseline_key.clone(),
|
||
(vec![vec![0.0f32; embed_dim]; num_layers], 0),
|
||
);
|
||
|
||
let mut total_seen = 0usize;
|
||
let mut total_used = 0usize;
|
||
|
||
for line in body.lines() {
|
||
if line.trim().is_empty() {
|
||
continue;
|
||
}
|
||
let row: Row = match serde_json::from_str(line) {
|
||
Ok(r) => r,
|
||
Err(e) => {
|
||
tracing::warn!("skipping unparseable row: {e}");
|
||
continue;
|
||
}
|
||
};
|
||
total_seen += 1;
|
||
let tag = match &row.emotion_tag {
|
||
Some(t) => normalize_tag(t),
|
||
None => continue,
|
||
};
|
||
let pool = match pools.get_mut(&tag) {
|
||
Some(p) => p,
|
||
None => continue,
|
||
};
|
||
if cli.max_per_pool > 0 && pool.1 >= cli.max_per_pool {
|
||
continue;
|
||
}
|
||
|
||
// Resolve WAV path.
|
||
let wav = if row.wav.is_absolute() {
|
||
row.wav.clone()
|
||
} else {
|
||
manifest_dir.join(&row.wav)
|
||
};
|
||
if !wav.exists() {
|
||
tracing::warn!("missing wav, skipping: {}", wav.display());
|
||
continue;
|
||
}
|
||
|
||
// Encode audio + build prompt. Mimi accumulates streaming state
|
||
// across encode() calls and the transformer's internal frame
|
||
// counter overflows its 8192-frame buffer after ~80 clips even
|
||
// with reset_state(). Pattern used elsewhere (commit 0568dd3):
|
||
// reset_state() per row, full reload() every 10.
|
||
let audio = audio_io::load_mono_24k(&wav)?;
|
||
let speaker = row.speaker.unwrap_or(0);
|
||
|
||
if total_used > 0 && total_used % 10 == 0 {
|
||
generator.mimi.reload()?;
|
||
} else {
|
||
generator.mimi.reset_state();
|
||
}
|
||
|
||
let captures = if module == "decoder" {
|
||
// Decoder mode: text-only segment for the prompt; Mimi-encode
|
||
// the audio separately to grab the middle frame's c0 as the
|
||
// teacher-forced target. This isolates the decoder's response
|
||
// to "given the text context and one audio token, what should
|
||
// codebooks 1..N-1 be?" — pure acoustic, no semantic prefix.
|
||
let segment = Segment::new_text(speaker, row.transcript.clone());
|
||
let prompt = build_prompt(
|
||
&[],
|
||
&segment,
|
||
&generator.model,
|
||
&mut generator.mimi,
|
||
&generator.tokenizer,
|
||
)?;
|
||
let codes = generator.mimi.encode(&audio)?;
|
||
let (_b, _cb, num_frames) = codes.dims3()?;
|
||
if num_frames == 0 {
|
||
tracing::warn!("zero-frame encode, skipping: {}", wav.display());
|
||
continue;
|
||
}
|
||
// Middle frame's c0 — past the audio onset, into steady state.
|
||
let mid = num_frames / 2;
|
||
let c0 = codes.narrow(2, mid, 1)?.narrow(1, 0, 1)?;
|
||
let target_c0: u32 = c0.to_dtype(DType::U32)?.flatten_all()?.to_vec1::<u32>()?[0];
|
||
generator.model.inner.capture_decoder_activations(
|
||
&prompt.tokens,
|
||
&prompt.mask,
|
||
target_c0,
|
||
)?
|
||
} else {
|
||
let segment = Segment::new(speaker, row.transcript.clone(), audio);
|
||
let prompt = build_prompt(
|
||
&[],
|
||
&segment,
|
||
&generator.model,
|
||
&mut generator.mimi,
|
||
&generator.tokenizer,
|
||
)?;
|
||
generator
|
||
.model
|
||
.inner
|
||
.capture_backbone_activations(&prompt.tokens, &prompt.mask)?
|
||
};
|
||
if captures.len() != num_layers {
|
||
anyhow::bail!(
|
||
"captured {} layers, expected {} — model layout mismatch?",
|
||
captures.len(),
|
||
num_layers
|
||
);
|
||
}
|
||
|
||
// Accumulate (move to CPU f32 for stable double-precision sum).
|
||
for (l, t) in captures.iter().enumerate() {
|
||
let v: Vec<f32> = t.to_dtype(DType::F32)?.to_vec1::<f32>()?;
|
||
if v.len() != embed_dim {
|
||
anyhow::bail!(
|
||
"layer {l} captured embed_dim {}, expected {embed_dim}",
|
||
v.len()
|
||
);
|
||
}
|
||
for (acc, vi) in pool.0[l].iter_mut().zip(v.iter()) {
|
||
*acc += *vi;
|
||
}
|
||
}
|
||
pool.1 += 1;
|
||
total_used += 1;
|
||
|
||
if total_used % 5 == 0 {
|
||
tracing::info!(
|
||
"progress: seen={total_seen} used={total_used} target={} baseline={}",
|
||
pools.get(&target_key).map(|p| p.1).unwrap_or(0),
|
||
pools.get(&baseline_key).map(|p| p.1).unwrap_or(0),
|
||
);
|
||
}
|
||
}
|
||
|
||
let n_target = pools.get(&target_key).map(|p| p.1).unwrap_or(0);
|
||
let n_base = pools.get(&baseline_key).map(|p| p.1).unwrap_or(0);
|
||
if n_target == 0 || n_base == 0 {
|
||
anyhow::bail!(
|
||
"need ≥1 row per pool: target='{}'={} baseline='{}'={}",
|
||
target_key,
|
||
n_target,
|
||
baseline_key,
|
||
n_base,
|
||
);
|
||
}
|
||
|
||
// Compute per-layer (mean(target) - mean(baseline)).
|
||
let mut diff_tensors: HashMap<String, Tensor> = HashMap::new();
|
||
for layer_idx in 0..num_layers {
|
||
let target_mean: Vec<f32> = pools[&target_key].0[layer_idx]
|
||
.iter()
|
||
.map(|v| v / n_target as f32)
|
||
.collect();
|
||
let base_mean: Vec<f32> = pools[&baseline_key].0[layer_idx]
|
||
.iter()
|
||
.map(|v| v / n_base as f32)
|
||
.collect();
|
||
let diff: Vec<f32> = target_mean
|
||
.iter()
|
||
.zip(base_mean.iter())
|
||
.map(|(t, b)| t - b)
|
||
.collect();
|
||
let t = Tensor::from_vec(diff, (embed_dim,), &Device::Cpu)?;
|
||
diff_tensors.insert(format!("layer_{layer_idx}_steering"), t);
|
||
}
|
||
|
||
if let Some(parent) = cli.out.parent() {
|
||
std::fs::create_dir_all(parent).ok();
|
||
}
|
||
candle_core::safetensors::save(&diff_tensors, &cli.out)?;
|
||
eprintln!("--- summary ---");
|
||
eprintln!("target '{target_key}': {n_target} rows");
|
||
eprintln!("baseline '{baseline_key}': {n_base} rows");
|
||
eprintln!("rows seen / used: {total_seen} / {total_used}");
|
||
eprintln!("wrote {} layers to {}", num_layers, cli.out.display());
|
||
Ok(())
|
||
}
|