mirror of
https://github.com/block/buzz.git
synced 2026-08-18 06:50:31 +02:00
fix(huddle): Pocket TTS quality — grouped chunking + FP32 model bundle (#2224)
Signed-off-by: npub1mprnacetjua2xx3p5eddmhxyk6wv929ymm5py8kd2xfxurxahspqqlgyta <d8473ee32b973aa31a21a65adddcc4b69cc2a8a4dee8121ecd51926e0cddbc02@sprout-oss.stage.blox.sqprod.co> Signed-off-by: Tyler Longwell <tlongwell@block.xyz> Co-authored-by: npub1mprnacetjua2xx3p5eddmhxyk6wv929ymm5py8kd2xfxurxahspqqlgyta <d8473ee32b973aa31a21a65adddcc4b69cc2a8a4dee8121ecd51926e0cddbc02@sprout-oss.stage.blox.sqprod.co> Co-authored-by: npub1qyvc0c5kl4gqv2fd97fsk46tu378sqgy35vc83rvgfwne90sel7s0ed67d <011987e296fd5006292d2f930b574be47c7801048d1983c46c425d3c95f0cffd@sprout-oss.stage.blox.sqprod.co> Co-authored-by: Tyler Longwell <tlongwell@block.xyz>
This commit is contained in:
Tyler
co-authored by
npub1mprnacetjua2xx3p5eddmhxyk6wv929ymm5py8kd2xfxurxahspqqlgyta
npub1qyvc0c5kl4gqv2fd97fsk46tu378sqgy35vc83rvgfwne90sel7s0ed67d
Tyler Longwell
npub1mprnacetjua2xx3p5eddmhxyk6wv929ymm5py8kd2xfxurxahspqqlgyta
npub1qyvc0c5kl4gqv2fd97fsk46tu378sqgy35vc83rvgfwne90sel7s0ed67d
Tyler Longwell
parent
47d7eb6982
commit
87276320af
@@ -0,0 +1,519 @@
|
||||
//! Reproducible blind Pocket TTS quality corpus generator.
|
||||
//!
|
||||
//! Renders Buzz's production prompt preparation and post-processing across:
|
||||
//! INT8/FP32 × per-sentence/grouped generation. The generated filenames are
|
||||
//! deterministically blinded; keep `key.json` away from listeners until their
|
||||
//! scoring sheet is complete.
|
||||
//!
|
||||
//! Usage:
|
||||
//! cargo run --release --example pocket_quality_ab -- \
|
||||
//! <int8-model-dir> <fp32-model-dir> <output-dir> [--idle-minutes N --only ITEM]
|
||||
//!
|
||||
//! The optional idle run intentionally creates one engine per condition, warms
|
||||
//! all four, sleeps once, and then makes each clip the first generation after
|
||||
//! dormancy. It requires `--only` because only the first synthesis after an
|
||||
//! uninterrupted idle is a valid post-idle observation. Run each 5/15-minute
|
||||
//! item as a separate process.
|
||||
|
||||
// Importing the production module also brings in runtime-only helpers that this
|
||||
// standalone corpus generator deliberately does not call.
|
||||
#![allow(dead_code)]
|
||||
|
||||
#[path = "../src/huddle/pocket.rs"]
|
||||
mod production_pocket;
|
||||
#[path = "../src/huddle/preprocessing.rs"]
|
||||
mod production_preprocessing;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use serde::Serialize;
|
||||
use sha2::{Digest, Sha256};
|
||||
use sherpa_onnx::{GenerationConfig, OfflineTts, OfflineTtsConfig, Wave};
|
||||
|
||||
use production_pocket::{prepare_pocket_prompt, SAMPLE_RATE};
|
||||
use production_preprocessing::{preprocess_for_tts, split_sentences};
|
||||
|
||||
const NUM_STEPS: i32 = 1;
|
||||
const SILENCE_SCALE: f32 = 1.0;
|
||||
const INTER_SENTENCE_SILENCE_SAMPLES: usize = SAMPLE_RATE as usize / 10;
|
||||
const LEAD_IN_SAMPLES: usize = SAMPLE_RATE as usize / 50;
|
||||
const FADE_OUT_SAMPLES: usize = SAMPLE_RATE as usize * 8 / 1000;
|
||||
const TARGET_RMS_DBFS: f32 = -23.0;
|
||||
const BLINDING_SEED: &str = "pocket-quality-2026-07-21-v1";
|
||||
|
||||
const CORPUS: &[CorpusItem] = &[
|
||||
CorpusItem { id: "short_one_word", kind: "short", text: "Yep." },
|
||||
CorpusItem { id: "short_four_words", kind: "short", text: "Sounds good to me." },
|
||||
CorpusItem {
|
||||
id: "multi_relay_review",
|
||||
kind: "multi-sentence",
|
||||
text: "I looked at the relay code this morning. The lease logic is solid. There's one race in the worker claim path, though. I'll write it up and send you a patch.",
|
||||
},
|
||||
CorpusItem {
|
||||
id: "multi_community_size",
|
||||
kind: "multi-sentence",
|
||||
text: "Great question. The answer is it depends on the community size. For small ones, keep it simple.",
|
||||
},
|
||||
CorpusItem {
|
||||
id: "mixed_agent_message",
|
||||
kind: "mixed",
|
||||
text: "That's 42 open PRs right now — mostly small. I'll triage them after lunch.",
|
||||
},
|
||||
];
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
struct CorpusItem {
|
||||
id: &'static str,
|
||||
kind: &'static str,
|
||||
text: &'static str,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
enum Precision {
|
||||
Int8,
|
||||
Fp32,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
enum Chunking {
|
||||
PerSentence,
|
||||
Grouped,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
struct Condition {
|
||||
precision: Precision,
|
||||
chunking: Chunking,
|
||||
}
|
||||
|
||||
const CONDITIONS: [Condition; 4] = [
|
||||
Condition {
|
||||
precision: Precision::Int8,
|
||||
chunking: Chunking::PerSentence,
|
||||
},
|
||||
Condition {
|
||||
precision: Precision::Int8,
|
||||
chunking: Chunking::Grouped,
|
||||
},
|
||||
Condition {
|
||||
precision: Precision::Fp32,
|
||||
chunking: Chunking::PerSentence,
|
||||
},
|
||||
Condition {
|
||||
precision: Precision::Fp32,
|
||||
chunking: Chunking::Grouped,
|
||||
},
|
||||
];
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct KeyFile {
|
||||
warning: &'static str,
|
||||
blinding_seed: &'static str,
|
||||
target_rms_dbfs: f32,
|
||||
items: Vec<KeyItem>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct KeyItem {
|
||||
id: String,
|
||||
kind: String,
|
||||
text: String,
|
||||
clips: Vec<KeyClip>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct KeyClip {
|
||||
file: String,
|
||||
precision: Precision,
|
||||
chunking: Chunking,
|
||||
cold_start: bool,
|
||||
idle_minutes: Option<u64>,
|
||||
synthesis_ms: u128,
|
||||
audio_seconds: f32,
|
||||
}
|
||||
|
||||
struct Voice {
|
||||
samples: Vec<f32>,
|
||||
sample_rate: i32,
|
||||
}
|
||||
|
||||
struct Engine {
|
||||
inner: OfflineTts,
|
||||
voice: Voice,
|
||||
}
|
||||
|
||||
fn main() -> Result<(), String> {
|
||||
let mut args = std::env::args().skip(1);
|
||||
let int8_dir = required_path(args.next(), "INT8 model directory")?;
|
||||
let fp32_dir = required_path(args.next(), "FP32 model directory")?;
|
||||
let output_dir = required_path(args.next(), "output directory")?;
|
||||
let mut idle_minutes = None;
|
||||
let mut only_item = None;
|
||||
while let Some(arg) = args.next() {
|
||||
match arg.as_str() {
|
||||
"--idle-minutes" => {
|
||||
idle_minutes = Some(
|
||||
args.next()
|
||||
.ok_or("--idle-minutes requires a value")?
|
||||
.parse::<u64>()
|
||||
.map_err(|e| format!("invalid idle minutes: {e}"))?,
|
||||
);
|
||||
}
|
||||
"--only" => only_item = Some(args.next().ok_or("--only requires an item ID")?),
|
||||
_ => return Err(format!("unknown argument: {arg}")),
|
||||
}
|
||||
}
|
||||
|
||||
if idle_minutes.is_some() && only_item.is_none() {
|
||||
return Err("--idle-minutes requires --only so every clip is first-after-idle".into());
|
||||
}
|
||||
if let Some(ref requested) = only_item {
|
||||
if !CORPUS.iter().any(|item| item.id == requested) {
|
||||
return Err(format!("unknown corpus item for --only: {requested}"));
|
||||
}
|
||||
}
|
||||
|
||||
validate_model_dir(&int8_dir, Precision::Int8)?;
|
||||
validate_model_dir(&fp32_dir, Precision::Fp32)?;
|
||||
fs::create_dir_all(&output_dir).map_err(|e| e.to_string())?;
|
||||
|
||||
let mut engines = Vec::with_capacity(CONDITIONS.len());
|
||||
for condition in CONDITIONS {
|
||||
let dir = match condition.precision {
|
||||
Precision::Int8 => &int8_dir,
|
||||
Precision::Fp32 => &fp32_dir,
|
||||
};
|
||||
let engine = load_engine(dir, condition.precision)?;
|
||||
// Production warms once before serving a real utterance. Cold cases use
|
||||
// separate fresh engines below and deliberately skip this call.
|
||||
synth_chunks(&engine, &["warmup".to_string()])?;
|
||||
engines.push(engine);
|
||||
}
|
||||
|
||||
if let Some(minutes) = idle_minutes {
|
||||
eprintln!("All four warmed engines idle for {minutes} minute(s)…");
|
||||
std::thread::sleep(Duration::from_secs(minutes * 60));
|
||||
}
|
||||
|
||||
let mut key_items = Vec::new();
|
||||
for item in CORPUS {
|
||||
if only_item
|
||||
.as_deref()
|
||||
.is_some_and(|requested| requested != item.id)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let preprocessed = preprocess_for_tts(item.text);
|
||||
let per_sentence: Vec<String> = split_sentences(&preprocessed)
|
||||
.into_iter()
|
||||
.filter(|s| !s.trim().is_empty())
|
||||
.collect();
|
||||
// These corpus texts are deliberately below the upstream ~50-token
|
||||
// grouping target, so grouped mode is one exact generate() call.
|
||||
let grouped = vec![per_sentence.join(" ")];
|
||||
let item_dir = output_dir.join(item.id);
|
||||
fs::create_dir_all(&item_dir).map_err(|e| e.to_string())?;
|
||||
let clip_order = blinded_order(item.id);
|
||||
let mut clips = Vec::new();
|
||||
|
||||
let mut rendered = Vec::new();
|
||||
for (condition_index, engine) in engines.iter().enumerate() {
|
||||
let condition = CONDITIONS[condition_index];
|
||||
let chunks = match condition.chunking {
|
||||
Chunking::PerSentence => &per_sentence,
|
||||
Chunking::Grouped => &grouped,
|
||||
};
|
||||
let started = Instant::now();
|
||||
let audio = synth_chunks(engine, chunks)?;
|
||||
rendered.push((
|
||||
condition_index,
|
||||
condition,
|
||||
audio,
|
||||
started.elapsed().as_millis(),
|
||||
));
|
||||
}
|
||||
loudness_match_item(&mut rendered);
|
||||
for (condition_index, condition, audio, synth_ms) in rendered {
|
||||
let clip_number = clip_order[condition_index] + 1;
|
||||
let file_name = format!("clip{clip_number}.wav");
|
||||
write_wav(&item_dir.join(&file_name), &audio)?;
|
||||
clips.push(KeyClip {
|
||||
file: format!("{}/{file_name}", item.id),
|
||||
precision: condition.precision,
|
||||
chunking: condition.chunking,
|
||||
cold_start: false,
|
||||
idle_minutes,
|
||||
synthesis_ms: synth_ms,
|
||||
audio_seconds: audio.len() as f32 / SAMPLE_RATE as f32,
|
||||
});
|
||||
}
|
||||
clips.sort_by(|a, b| a.file.cmp(&b.file));
|
||||
key_items.push(KeyItem {
|
||||
id: item.id.to_string(),
|
||||
kind: item.kind.to_string(),
|
||||
text: item.text.to_string(),
|
||||
clips,
|
||||
});
|
||||
}
|
||||
|
||||
// Explicit fresh-engine cold-start clips for the two highest-signal texts.
|
||||
// Idle runs intentionally omit them: they happen after the post-idle clips
|
||||
// and add no valid idle observation.
|
||||
for item in if idle_minutes.is_none() { CORPUS } else { &[] } {
|
||||
if !matches!(item.id, "short_one_word" | "multi_relay_review") {
|
||||
continue;
|
||||
}
|
||||
if only_item
|
||||
.as_deref()
|
||||
.is_some_and(|requested| requested != item.id)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let cold_id = format!("cold_{}", item.id);
|
||||
let preprocessed = preprocess_for_tts(item.text);
|
||||
let sentences: Vec<String> = split_sentences(&preprocessed)
|
||||
.into_iter()
|
||||
.filter(|s| !s.trim().is_empty())
|
||||
.collect();
|
||||
let grouped = vec![sentences.join(" ")];
|
||||
let item_dir = output_dir.join(&cold_id);
|
||||
fs::create_dir_all(&item_dir).map_err(|e| e.to_string())?;
|
||||
let clip_order = blinded_order(&cold_id);
|
||||
let mut clips = Vec::new();
|
||||
let mut rendered = Vec::new();
|
||||
for (condition_index, condition) in CONDITIONS.iter().copied().enumerate() {
|
||||
let dir = match condition.precision {
|
||||
Precision::Int8 => &int8_dir,
|
||||
Precision::Fp32 => &fp32_dir,
|
||||
};
|
||||
let engine = load_engine(dir, condition.precision)?;
|
||||
let chunks = match condition.chunking {
|
||||
Chunking::PerSentence => &sentences,
|
||||
Chunking::Grouped => &grouped,
|
||||
};
|
||||
let started = Instant::now();
|
||||
let audio = synth_chunks(&engine, chunks)?;
|
||||
rendered.push((
|
||||
condition_index,
|
||||
condition,
|
||||
audio,
|
||||
started.elapsed().as_millis(),
|
||||
));
|
||||
}
|
||||
loudness_match_item(&mut rendered);
|
||||
for (condition_index, condition, audio, synth_ms) in rendered {
|
||||
let clip_number = clip_order[condition_index] + 1;
|
||||
let file_name = format!("clip{clip_number}.wav");
|
||||
write_wav(&item_dir.join(&file_name), &audio)?;
|
||||
clips.push(KeyClip {
|
||||
file: format!("{cold_id}/{file_name}"),
|
||||
precision: condition.precision,
|
||||
chunking: condition.chunking,
|
||||
cold_start: true,
|
||||
idle_minutes: None,
|
||||
synthesis_ms: synth_ms,
|
||||
audio_seconds: audio.len() as f32 / SAMPLE_RATE as f32,
|
||||
});
|
||||
}
|
||||
clips.sort_by(|a, b| a.file.cmp(&b.file));
|
||||
key_items.push(KeyItem {
|
||||
id: cold_id,
|
||||
kind: "cold-start".to_string(),
|
||||
text: item.text.to_string(),
|
||||
clips,
|
||||
});
|
||||
}
|
||||
|
||||
let key = KeyFile {
|
||||
warning: "DO NOT OPEN UNTIL LISTENING SCORES ARE FINAL",
|
||||
blinding_seed: BLINDING_SEED,
|
||||
target_rms_dbfs: TARGET_RMS_DBFS,
|
||||
items: key_items,
|
||||
};
|
||||
fs::write(
|
||||
output_dir.join("key.json"),
|
||||
serde_json::to_vec_pretty(&key).map_err(|e| e.to_string())?,
|
||||
)
|
||||
.map_err(|e| e.to_string())?;
|
||||
write_scoring_sheet(&output_dir, &key)?;
|
||||
println!("Wrote blind corpus to {}", output_dir.display());
|
||||
println!("Give listeners the WAV folders and SCORING.md; withhold key.json.");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn required_path(value: Option<String>, label: &str) -> Result<PathBuf, String> {
|
||||
value
|
||||
.map(PathBuf::from)
|
||||
.ok_or_else(|| format!("missing {label}"))
|
||||
}
|
||||
|
||||
fn model_file(precision: Precision, base: &str) -> String {
|
||||
match precision {
|
||||
Precision::Int8 => format!("{base}.int8.onnx"),
|
||||
Precision::Fp32 => format!("{base}.onnx"),
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_model_dir(dir: &Path, precision: Precision) -> Result<(), String> {
|
||||
for file in [
|
||||
model_file(precision, "lm_main"),
|
||||
model_file(precision, "lm_flow"),
|
||||
"encoder.onnx".into(),
|
||||
model_file(precision, "decoder"),
|
||||
"text_conditioner.onnx".into(),
|
||||
"vocab.json".into(),
|
||||
"token_scores.json".into(),
|
||||
"reference_sample.wav".into(),
|
||||
] {
|
||||
if !dir.join(&file).is_file() {
|
||||
return Err(format!("missing {}", dir.join(file).display()));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn load_engine(dir: &Path, precision: Precision) -> Result<Engine, String> {
|
||||
let p = |name: &str| dir.join(name).to_string_lossy().into_owned();
|
||||
let mut cfg = OfflineTtsConfig::default();
|
||||
cfg.model.pocket.lm_main = Some(p(&model_file(precision, "lm_main")));
|
||||
cfg.model.pocket.lm_flow = Some(p(&model_file(precision, "lm_flow")));
|
||||
cfg.model.pocket.encoder = Some(p("encoder.onnx"));
|
||||
cfg.model.pocket.decoder = Some(p(&model_file(precision, "decoder")));
|
||||
cfg.model.pocket.text_conditioner = Some(p("text_conditioner.onnx"));
|
||||
cfg.model.pocket.vocab_json = Some(p("vocab.json"));
|
||||
cfg.model.pocket.token_scores_json = Some(p("token_scores.json"));
|
||||
cfg.model.pocket.voice_embedding_cache_capacity = 16;
|
||||
cfg.model.num_threads = 1;
|
||||
cfg.model.debug = false;
|
||||
let inner =
|
||||
OfflineTts::create(&cfg).ok_or_else(|| format!("failed to create {precision:?} engine"))?;
|
||||
let wave =
|
||||
Wave::read(&p("reference_sample.wav")).ok_or("failed to read reference_sample.wav")?;
|
||||
Ok(Engine {
|
||||
inner,
|
||||
voice: Voice {
|
||||
samples: wave.samples().to_vec(),
|
||||
sample_rate: wave.sample_rate(),
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
fn synth_chunks(engine: &Engine, chunks: &[String]) -> Result<Vec<f32>, String> {
|
||||
let mut out = Vec::new();
|
||||
for chunk in chunks {
|
||||
let prepared = prepare_pocket_prompt(chunk).ok_or("empty prepared prompt")?;
|
||||
let extra = prepared.max_frames.map(|max_frames| {
|
||||
HashMap::from([(
|
||||
"max_frames".to_string(),
|
||||
serde_json::Value::from(max_frames),
|
||||
)])
|
||||
});
|
||||
let cfg = GenerationConfig {
|
||||
num_steps: NUM_STEPS,
|
||||
silence_scale: SILENCE_SCALE,
|
||||
reference_audio: Some(engine.voice.samples.clone()),
|
||||
reference_sample_rate: engine.voice.sample_rate,
|
||||
extra,
|
||||
..Default::default()
|
||||
};
|
||||
let audio = engine
|
||||
.inner
|
||||
.generate_with_config(&prepared.text, &cfg, None::<fn(&[f32], f32) -> bool>)
|
||||
.ok_or_else(|| format!("synthesis failed for {chunk:?}"))?;
|
||||
let mut samples: Vec<f32> = audio.samples().iter().map(|s| s.clamp(-1.0, 1.0)).collect();
|
||||
apply_fade_out(&mut samples);
|
||||
out.extend(std::iter::repeat_n(0.0, LEAD_IN_SAMPLES));
|
||||
out.extend(samples);
|
||||
out.extend(std::iter::repeat_n(
|
||||
0.0,
|
||||
INTER_SENTENCE_SILENCE_SAMPLES - LEAD_IN_SAMPLES,
|
||||
));
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
fn apply_fade_out(samples: &mut [f32]) {
|
||||
let fade = FADE_OUT_SAMPLES.min(samples.len() / 2);
|
||||
for i in 0..fade {
|
||||
samples[samples.len() - 1 - i] *= i as f32 / fade as f32;
|
||||
}
|
||||
}
|
||||
|
||||
fn active_rms(samples: &[f32]) -> Option<f32> {
|
||||
let (sum_squares, count) = samples
|
||||
.iter()
|
||||
.filter(|sample| sample.abs() > 1.0e-4)
|
||||
.fold((0.0_f32, 0_usize), |(sum, count), sample| {
|
||||
(sum + sample * sample, count + 1)
|
||||
});
|
||||
(count > 0).then(|| (sum_squares / count as f32).sqrt())
|
||||
}
|
||||
|
||||
/// Attenuate every clip in one comparison set to the quietest active-speech RMS.
|
||||
/// This removes the louder-is-better confound without normalizing dynamics or
|
||||
/// claiming standards-compliant integrated LUFS. The dBFS value is a ceiling.
|
||||
fn loudness_match_item(rendered: &mut [(usize, Condition, Vec<f32>, u128)]) {
|
||||
let ceiling = 10.0_f32.powf(TARGET_RMS_DBFS / 20.0);
|
||||
let target = rendered
|
||||
.iter()
|
||||
.filter_map(|(_, _, samples, _)| active_rms(samples))
|
||||
.fold(ceiling, f32::min);
|
||||
for (_, _, samples, _) in rendered {
|
||||
let Some(rms) = active_rms(samples) else {
|
||||
continue;
|
||||
};
|
||||
let gain = (target / rms).min(1.0);
|
||||
for sample in samples {
|
||||
*sample *= gain;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn blinded_order(item_id: &str) -> [usize; 4] {
|
||||
let mut keyed: Vec<(usize, Vec<u8>)> = (0..4)
|
||||
.map(|index| {
|
||||
let digest = Sha256::digest(format!("{BLINDING_SEED}:{item_id}:{index}"));
|
||||
(index, digest.to_vec())
|
||||
})
|
||||
.collect();
|
||||
keyed.sort_by(|a, b| a.1.cmp(&b.1));
|
||||
let mut condition_to_clip = [0; 4];
|
||||
for (clip, (condition, _)) in keyed.into_iter().enumerate() {
|
||||
condition_to_clip[condition] = clip;
|
||||
}
|
||||
condition_to_clip
|
||||
}
|
||||
|
||||
fn write_wav(path: &Path, samples: &[f32]) -> Result<(), String> {
|
||||
let path = path
|
||||
.to_str()
|
||||
.ok_or_else(|| format!("non-UTF8 path: {}", path.display()))?;
|
||||
if sherpa_onnx::write(path, samples, SAMPLE_RATE as i32) {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(format!("failed to write {path}"))
|
||||
}
|
||||
}
|
||||
|
||||
fn write_scoring_sheet(output_dir: &Path, key: &KeyFile) -> Result<(), String> {
|
||||
let mut sheet = String::from("# Pocket TTS blind listening sheet\n\nDo not open `key.json` until this sheet is complete. Rank best to worst; ties are allowed.\n\n");
|
||||
for item in &key.items {
|
||||
sheet.push_str(&format!(
|
||||
"## {} ({})\n\n> {}\n\n",
|
||||
item.id, item.kind, item.text
|
||||
));
|
||||
sheet.push_str("Rank: `____ > ____ > ____ > ____`\n\n| Clip | seam | onset | garble | robotic | timbre | truncate | note |\n|---|---|---|---|---|---|---|---|\n");
|
||||
for clip in 1..=4 {
|
||||
sheet.push_str(&format!(
|
||||
"| clip{clip} | [ ] | [ ] | [ ] | [ ] | [ ] | [ ] | |\n"
|
||||
));
|
||||
}
|
||||
sheet.push('\n');
|
||||
}
|
||||
fs::write(output_dir.join("SCORING.md"), sheet).map_err(|e| e.to_string())
|
||||
}
|
||||
@@ -38,12 +38,19 @@ use sha2::{Digest, Sha256};
|
||||
/// Computed from a known-good download. Update when upgrading model versions.
|
||||
const STT_ARCHIVE_SHA256: &str = "17f945007b52ccd8b7200ffc7c5652e9e8e961dfdf479cefcabd06cf5703630b";
|
||||
|
||||
/// HuggingFace base URL for the sherpa-onnx Pocket TTS int8 repackage.
|
||||
/// HuggingFace base URL for the sherpa-onnx Pocket TTS fp32 repackage.
|
||||
///
|
||||
/// Pinned to commit e715955cf50d18d919d37231513c0e914b83661a
|
||||
/// Pinned to commit 96d1e53ce3311ca6c2c6a35e2062d36b4cec6fa3
|
||||
/// (2026-02-10) for reproducible downloads.
|
||||
///
|
||||
/// fp32 (not int8): a direct same-runtime A/B (k2-fsa/sherpa-onnx#3172)
|
||||
/// found the ONNX int8 quantization audibly degraded Pocket TTS output and
|
||||
/// that fp32 "significantly improved quality even at 1 step". The runtime
|
||||
/// bundle grows from ~189 MB to ~473 MB; encoder, text conditioner, both
|
||||
/// JSON tables, and LICENSE are byte-identical between the two repos — only
|
||||
/// the three quantized sessions (lm_main, lm_flow, decoder) change.
|
||||
const POCKET_HF_BASE: &str =
|
||||
"https://huggingface.co/csukuangfj2/sherpa-onnx-pocket-tts-int8-2026-01-26/resolve/e715955cf50d18d919d37231513c0e914b83661a";
|
||||
"https://huggingface.co/csukuangfj2/sherpa-onnx-pocket-tts-2026-01-26/resolve/96d1e53ce3311ca6c2c6a35e2062d36b4cec6fa3";
|
||||
|
||||
/// Reference voice WAV: "Mary (f, conversation)" from the Kyutai TTS demo
|
||||
/// voice set — VCTK speaker p333, ai-coustics-enhanced. Pinned to
|
||||
@@ -61,10 +68,10 @@ const POCKET_REFERENCE_WAV_URL: &str =
|
||||
/// Computed from known-good pinned downloads. Update when upgrading model versions.
|
||||
#[rustfmt::skip]
|
||||
const TTS_FILE_HASHES: &[(&str, &str)] = &[
|
||||
("decoder.int8.onnx", "12b0857402d31aead94df19d6783b4350d1f740e811f3a3202c70ad89ae11eea"),
|
||||
("decoder.onnx", "f267880fde6c58b17b0a8f3647eaf8dcfad321f833f32d583ebc2fb2d1a15f10"),
|
||||
("encoder.onnx", "e8f2f6d301ffb96e398b138a7dc6d3038622d236044636b73d920bab85890260"),
|
||||
("lm_flow.int8.onnx", "8d627d235c44a597da908e1085ebe241cbbe358964c502c5a5063d18851a5529"),
|
||||
("lm_main.int8.onnx", "bfc0c7e7e3d72864fa3bb2ee499f62f21ddc1474b885f5f3ca570f8be73e787e"),
|
||||
("lm_flow.onnx", "79c013a554a54e63319c33c0cc8830cbbedc9b7e448ae7e26f7923ae11f9873e"),
|
||||
("lm_main.onnx", "255d1a9263c5abdf36034abfc19c11d21cc5f40f0f87d8361288e972cbd5c578"),
|
||||
("text_conditioner.onnx", "0b84e837d7bfaf2c896627b03e3f080320309f37f4fc7df7698c644f7ba5e6b1"),
|
||||
("vocab.json", "6fb646346cf931016f70c4921aab0900ce7a304b893cb02135c74e294abfea01"),
|
||||
("token_scores.json", "5be2f278caf9b9800741f0fd82bff677f4943ec764c356f907213434b622d958"),
|
||||
@@ -91,7 +98,9 @@ const STT_MODEL_VERSION: &str = "2";
|
||||
/// from kyutai/tts-voices. The hash mismatch on `reference_sample.wav` would
|
||||
/// fail readiness on its own, but the manifest bump makes the re-download
|
||||
/// reason explicit and skips the failing-then-re-fetching transient state.
|
||||
const TTS_MODEL_VERSION: &str = "2";
|
||||
/// Bumped "2" → "3" for the int8 → fp32 model swap (see `POCKET_HF_BASE`):
|
||||
/// existing int8 installs must re-download the suffixless fp32 sessions.
|
||||
const TTS_MODEL_VERSION: &str = "3";
|
||||
|
||||
/// Filename for the version manifest written alongside model files.
|
||||
const MANIFEST_FILENAME: &str = ".buzz-model-manifest";
|
||||
@@ -101,8 +110,9 @@ const MANIFEST_FILENAME: &str = ".buzz-model-manifest";
|
||||
/// Maximum expected STT archive size (200 MB — actual is ~100 MB).
|
||||
const MAX_STT_DOWNLOAD_BYTES: u64 = 200 * 1024 * 1024;
|
||||
|
||||
/// Maximum expected Pocket TTS file size (200 MB per file — largest is ~73 MB).
|
||||
const MAX_TTS_FILE_BYTES: u64 = 200 * 1024 * 1024;
|
||||
/// Maximum expected Pocket TTS file size (400 MB per file — largest is
|
||||
/// `lm_main.onnx` at ~303 MB fp32).
|
||||
const MAX_TTS_FILE_BYTES: u64 = 400 * 1024 * 1024;
|
||||
|
||||
/// NVIDIA Parakeet TDT-CTC 110M (English, int8) — packaged for sherpa-onnx by
|
||||
/// k2-fsa. Single ONNX file (CTC head) + tokens.txt. Avg WER ~7.5% across
|
||||
@@ -173,7 +183,7 @@ Mimi neural codec by Kyutai is bundled as part of the model.
|
||||
|
||||
ONNX export by KevinAHM: https://huggingface.co/KevinAHM/pocket-tts-onnx
|
||||
Sherpa-onnx repackage by csukuangfj / k2-fsa:
|
||||
https://huggingface.co/csukuangfj2/sherpa-onnx-pocket-tts-int8-2026-01-26
|
||||
https://huggingface.co/csukuangfj2/sherpa-onnx-pocket-tts-2026-01-26
|
||||
|
||||
Bundled reference voice (reference_sample.wav):
|
||||
\"Mary (f, conversation)\" preset from the Kyutai TTS demo voice catalogue
|
||||
@@ -193,10 +203,10 @@ license text for full warranty disclaimer.
|
||||
|
||||
/// All files that must be present for Pocket TTS to be considered ready.
|
||||
const TTS_EXPECTED_FILES: &[&str] = &[
|
||||
"decoder.int8.onnx",
|
||||
"decoder.onnx",
|
||||
"encoder.onnx",
|
||||
"lm_flow.int8.onnx",
|
||||
"lm_main.int8.onnx",
|
||||
"lm_flow.onnx",
|
||||
"lm_main.onnx",
|
||||
"text_conditioner.onnx",
|
||||
"vocab.json",
|
||||
"token_scores.json",
|
||||
@@ -759,10 +769,10 @@ impl ModelManager {
|
||||
fresh_temp_dir(&temp_dir).await?;
|
||||
|
||||
let model_files = [
|
||||
"decoder.int8.onnx",
|
||||
"decoder.onnx",
|
||||
"encoder.onnx",
|
||||
"lm_flow.int8.onnx",
|
||||
"lm_main.int8.onnx",
|
||||
"lm_flow.onnx",
|
||||
"lm_main.onnx",
|
||||
"text_conditioner.onnx",
|
||||
"vocab.json",
|
||||
"token_scores.json",
|
||||
|
||||
@@ -1,10 +1,15 @@
|
||||
//! Pocket TTS engine wrapper around sherpa-onnx's `OfflineTts`.
|
||||
//!
|
||||
//! Pocket TTS is a small (~189 MB int8 ONNX) zero-shot voice-cloning TTS
|
||||
//! Pocket TTS is a small (~473 MB fp32 ONNX) zero-shot voice-cloning TTS
|
||||
//! model from Kyutai. It runs quickly on CPU via sherpa-onnx, replacing the
|
||||
//! previous Kokoro-82M engine that also required an espeak-free but
|
||||
//! lexicon-heavy G2P pipeline (Misaki + CMUdict).
|
||||
//!
|
||||
//! Full-precision fp32 sessions, not the ~189 MB int8 quantization we
|
||||
//! originally shipped: a direct same-runtime A/B (k2-fsa/sherpa-onnx#3172)
|
||||
//! found the int8 ONNX export audibly degraded output quality, and fp32
|
||||
//! "significantly improved quality even at 1 step".
|
||||
//!
|
||||
//! ## Attribution
|
||||
//!
|
||||
//! - **Model**: Kyutai *Pocket TTS* — Charles, Roebel, et al., 2026.
|
||||
@@ -14,7 +19,7 @@
|
||||
//! - **ONNX export**: KevinAHM —
|
||||
//! <https://huggingface.co/KevinAHM/pocket-tts-onnx>. CC-BY-4.0.
|
||||
//! - **sherpa-onnx repackage**: csukuangfj / k2-fsa —
|
||||
//! <https://huggingface.co/csukuangfj2/sherpa-onnx-pocket-tts-int8-2026-01-26>.
|
||||
//! <https://huggingface.co/csukuangfj2/sherpa-onnx-pocket-tts-2026-01-26>.
|
||||
//! Repackages KevinAHM's export with the file layout sherpa-onnx's
|
||||
//! `OfflineTtsPocketModelConfig` expects. CC-BY-4.0.
|
||||
//! - **Reference voice WAV** (`reference_sample.wav`): the "Mary
|
||||
@@ -37,12 +42,16 @@
|
||||
//! - `VOICE_FILE_EXT: &str` — extension for per-voice files on disk.
|
||||
//! - `load_text_to_speech(model_dir)` → `Result<Engine, String>`
|
||||
//! - `load_voice_style(path)` → `Result<VoiceStyle, String>`
|
||||
//! - `Engine::synth_chunk(&self, text, lang, &VoiceStyle, steps, speed)`
|
||||
//! - `Engine::synth_chunk(&self, text, lang, &VoiceStyle, steps)`
|
||||
//! → `Result<Vec<f32>, String>`
|
||||
//!
|
||||
//! `lang` and `steps` are accepted for API compatibility with the previous
|
||||
//! Kokoro engine but are unused — Pocket TTS does its own language ID from
|
||||
//! the input text and is not a diffusion model (consistency LM, one step).
|
||||
//! There is no speed knob: sherpa-onnx's `GenerationConfig.speed` is only
|
||||
//! read by some model families (vits), never by the Pocket impl
|
||||
//! (`offline-tts-pocket-impl.h` — zero references), and upstream pocket-tts
|
||||
//! has no speed parameter either.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
@@ -143,10 +152,10 @@ const SHERPA_ONNX_FRAMES_AFTER_EOS_DEFAULT: i32 = 3;
|
||||
|
||||
// ── ONNX file names (five Pocket TTS sessions plus two JSON tables) ───────────
|
||||
|
||||
const FILE_LM_MAIN: &str = "lm_main.int8.onnx";
|
||||
const FILE_LM_FLOW: &str = "lm_flow.int8.onnx";
|
||||
const FILE_LM_MAIN: &str = "lm_main.onnx";
|
||||
const FILE_LM_FLOW: &str = "lm_flow.onnx";
|
||||
const FILE_ENCODER: &str = "encoder.onnx";
|
||||
const FILE_DECODER: &str = "decoder.int8.onnx";
|
||||
const FILE_DECODER: &str = "decoder.onnx";
|
||||
const FILE_TEXT_COND: &str = "text_conditioner.onnx";
|
||||
const FILE_VOCAB: &str = "vocab.json";
|
||||
const FILE_TOKEN_SCORES: &str = "token_scores.json";
|
||||
@@ -385,7 +394,6 @@ impl PocketTts {
|
||||
_lang: &str,
|
||||
style: &VoiceStyle,
|
||||
_steps: usize,
|
||||
speed: f32,
|
||||
) -> Result<Vec<f32>, String> {
|
||||
// Mirror upstream pocket-tts prompt prep — without this short or
|
||||
// unpunctuated inputs can cause the LM's EOS logit to never trip,
|
||||
@@ -405,12 +413,13 @@ impl PocketTts {
|
||||
let extra = build_generation_extra(&prepared);
|
||||
|
||||
let cfg = GenerationConfig {
|
||||
speed,
|
||||
num_steps: SYNTH_NUM_STEPS,
|
||||
silence_scale: SYNTH_SILENCE_SCALE,
|
||||
reference_audio: Some(style.samples.clone()),
|
||||
reference_sample_rate: style.sample_rate,
|
||||
extra,
|
||||
// `speed` stays at its default: the Pocket impl never reads it
|
||||
// (see the engine-contract note in the module docs).
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
|
||||
@@ -69,9 +69,6 @@ const MONITOR_TICK: Duration = Duration::from_millis(10);
|
||||
/// Pocket TTS is a one-step consistency model, not diffusion. Kept for API compat.
|
||||
const SYNTH_STEPS: usize = 1;
|
||||
|
||||
/// Synthesis speed multiplier. Slightly faster than natural speech.
|
||||
const SYNTH_SPEED: f32 = 1.05;
|
||||
|
||||
/// Fade-out length in samples (8 ms at 24 kHz ≈ 192 samples).
|
||||
///
|
||||
/// Applied only at the *end* of each synthesised sentence to eliminate the
|
||||
@@ -94,10 +91,24 @@ const FADE_OUT_SAMPLES: usize = (SAMPLE_RATE as f64 * 0.008) as usize;
|
||||
/// existing inter-sentence pause, so it does not lengthen multi-sentence gaps.
|
||||
const SENTENCE_LEAD_IN_SAMPLES: usize = (SAMPLE_RATE as f64 * 0.020) as usize;
|
||||
|
||||
/// Sentence-by-sentence synthesis — keeps first-sentence latency low and lets
|
||||
/// playback of sentence N overlap with synthesis of sentence N+1 (see the
|
||||
/// lookahead pipelining note in the module doc-comment above).
|
||||
const BATCH_SIZE: usize = 1;
|
||||
/// Approximate character budget for one synthesis chunk.
|
||||
///
|
||||
/// Upstream pocket-tts groups sentences into chunks of up to
|
||||
/// `MAX_TOKEN_PER_CHUNK = 50` tokenizer tokens (`default_parameters.py`) —
|
||||
/// typically multi-sentence chunks — because every `generate()` call is an
|
||||
/// independent generation with a cold FlowLM start, and each chunk boundary
|
||||
/// is an exposed prosody seam (kyutai-labs/pocket-tts #151; the Kyutai team
|
||||
/// names chunk stitching as the reliability lever). Our previous
|
||||
/// sentence-per-call path created ~2–4× more seams than upstream.
|
||||
///
|
||||
/// We don't ship the SentencePiece tokenizer, so 50 tokens is approximated
|
||||
/// with a character budget. The bundled 4k-entry vocab averages ~4 chars per
|
||||
/// token, but usage-weighted English text leans on short common tokens, so
|
||||
/// the effective ratio is ~2–4 chars/token and 200 chars ≈ 60–100 tokens —
|
||||
/// modestly above upstream's 50, deliberately: erring large means fewer
|
||||
/// seams, and even ~100 tokens is far below the model's 500-LM-step (~40 s)
|
||||
/// ceiling. Do not shrink this budget to chase an exact 50-token match.
|
||||
const MAX_CHUNK_CHARS: usize = 200;
|
||||
|
||||
/// Silence inserted between sentences by the TTS pipeline (seconds).
|
||||
/// Injected as a silent buffer between each synthesized sentence chunk.
|
||||
@@ -275,7 +286,7 @@ fn tts_worker(
|
||||
// and discard the output so the first real utterance runs at warm-session speed.
|
||||
{
|
||||
let t = std::time::Instant::now();
|
||||
match engine.synth_chunk("warmup", "en", &style, SYNTH_STEPS, SYNTH_SPEED) {
|
||||
match engine.synth_chunk("warmup", "en", &style, SYNTH_STEPS) {
|
||||
Ok(_) => eprintln!(
|
||||
"buzz-desktop: TTS warmup completed in {:.0}ms",
|
||||
t.elapsed().as_millis()
|
||||
@@ -480,15 +491,20 @@ fn tts_worker(
|
||||
continue;
|
||||
}
|
||||
|
||||
// Split into sentences. Each sentence is synthesized individually and
|
||||
// appended to the Player immediately — synthesis of sentence N+1 overlaps
|
||||
// with playback of sentence N (lookahead pipelining).
|
||||
// Split into sentences, then group into synthesis chunks: the first
|
||||
// sentence stays alone (fast time-to-first-audio), the rest pack
|
||||
// greedily up to MAX_CHUNK_CHARS. Each chunk is one `generate()`
|
||||
// call; playback of chunk N overlaps synthesis of chunk N+1
|
||||
// (lookahead pipelining). Grouping matches upstream's ~50-token
|
||||
// chunking and halves the exposed prosody seams on multi-sentence
|
||||
// replies — see MAX_CHUNK_CHARS.
|
||||
let sentences: Vec<String> = split_sentences(&text)
|
||||
.into_iter()
|
||||
.filter(|s| !s.trim().is_empty())
|
||||
.collect();
|
||||
let chunks = group_sentences_into_chunks(&sentences, MAX_CHUNK_CHARS);
|
||||
|
||||
for sentence in &sentences {
|
||||
for chunk in &chunks {
|
||||
if handle_cancel_or_shutdown(
|
||||
&cancel,
|
||||
&shutdown,
|
||||
@@ -500,12 +516,12 @@ fn tts_worker(
|
||||
break;
|
||||
}
|
||||
|
||||
let text = sentence.trim();
|
||||
let text = chunk.trim();
|
||||
if text.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
match engine.synth_chunk(text, "en", &style, SYNTH_STEPS, SYNTH_SPEED) {
|
||||
match engine.synth_chunk(text, "en", &style, SYNTH_STEPS) {
|
||||
Ok(samples) if !samples.is_empty() => {
|
||||
let mut audio = clamp_to_full_scale(samples);
|
||||
// Fade-out only — fading-in would attenuate the consonant
|
||||
@@ -711,13 +727,50 @@ fn build_sentence_append_buffer(
|
||||
buf
|
||||
}
|
||||
|
||||
/// Group sentences into synthesis chunks.
|
||||
///
|
||||
/// The first sentence always stands alone — it is what the listener hears
|
||||
/// first, and synthesizing it by itself keeps time-to-first-audio at the
|
||||
/// single-sentence cost. Subsequent sentences pack greedily: a sentence
|
||||
/// joins the current chunk while the combined length stays within
|
||||
/// `max_chars`; otherwise it starts a new chunk. A single sentence longer
|
||||
/// than `max_chars` becomes its own chunk unsplit — Pocket TTS handles long
|
||||
/// single sentences fine (the ceiling is the 500-LM-step default), it's the
|
||||
/// *seams* we're minimizing.
|
||||
///
|
||||
/// Sentences within a chunk are joined with a single space; sentence-ending
|
||||
/// punctuation is preserved by `split_sentences`, so the model sees natural
|
||||
/// multi-sentence prose — the same shape upstream's ~50-token chunker feeds it.
|
||||
fn group_sentences_into_chunks(sentences: &[String], max_chars: usize) -> Vec<String> {
|
||||
let mut chunks: Vec<String> = Vec::new();
|
||||
for (i, sentence) in sentences.iter().enumerate() {
|
||||
let sentence = sentence.trim();
|
||||
if sentence.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if i == 0 || chunks.is_empty() {
|
||||
chunks.push(sentence.to_string());
|
||||
continue;
|
||||
}
|
||||
// Never merge into the first chunk — it's the latency-critical one.
|
||||
let can_merge = chunks.len() > 1
|
||||
&& chunks
|
||||
.last()
|
||||
.is_some_and(|c| c.len() + 1 + sentence.len() <= max_chars);
|
||||
if can_merge {
|
||||
let last = chunks.last_mut().expect("non-empty checked above");
|
||||
last.push(' ');
|
||||
last.push_str(sentence);
|
||||
} else {
|
||||
chunks.push(sentence.to_string());
|
||||
}
|
||||
}
|
||||
chunks
|
||||
}
|
||||
|
||||
// drain_until_shutdown lives in super (huddle/mod.rs) — shared with stt.rs.
|
||||
use super::drain_until_shutdown;
|
||||
|
||||
// BATCH_SIZE is used implicitly (one sentence per iteration). Suppress dead_code
|
||||
// lint since it documents the design intent.
|
||||
const _: () = assert!(BATCH_SIZE == 1);
|
||||
|
||||
// ── Tests ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -915,3 +915,80 @@ fn clamp_to_full_scale_empty_buffer() {
|
||||
let out = clamp_to_full_scale(Vec::new());
|
||||
assert!(out.is_empty());
|
||||
}
|
||||
|
||||
// ── group_sentences_into_chunks tests ─────────────────────────────────────
|
||||
|
||||
fn s(v: &[&str]) -> Vec<String> {
|
||||
v.iter().map(|x| x.to_string()).collect()
|
||||
}
|
||||
|
||||
/// The first sentence always stands alone — it bounds time-to-first-audio.
|
||||
/// Even when the whole message would fit in one chunk, sentence one must
|
||||
/// not wait on synthesis of the rest.
|
||||
#[test]
|
||||
fn chunk_grouping_first_sentence_is_always_alone() {
|
||||
let chunks = group_sentences_into_chunks(&s(&["Hi there.", "Short.", "Tiny."]), 200);
|
||||
assert_eq!(chunks[0], "Hi there.");
|
||||
assert_eq!(chunks.len(), 2);
|
||||
assert_eq!(chunks[1], "Short. Tiny.");
|
||||
}
|
||||
|
||||
/// Sentences after the first pack greedily up to the char budget, then
|
||||
/// spill into a new chunk. Fewer generate() calls = fewer prosody seams.
|
||||
#[test]
|
||||
fn chunk_grouping_packs_up_to_budget_then_spills() {
|
||||
let a = "A".repeat(50) + ".";
|
||||
let b = "B".repeat(50) + ".";
|
||||
let c = "C".repeat(50) + ".";
|
||||
let d = "D".repeat(50) + ".";
|
||||
// Budget of 110: b+c fits (51+1+51 = 103), adding d (103+1+51) does not.
|
||||
let chunks = group_sentences_into_chunks(&s(&[&a, &b, &c, &d]), 110);
|
||||
assert_eq!(chunks.len(), 3, "chunks: {chunks:?}");
|
||||
assert_eq!(chunks[0], a);
|
||||
assert_eq!(chunks[1], format!("{b} {c}"));
|
||||
assert_eq!(chunks[2], d);
|
||||
}
|
||||
|
||||
/// A single sentence longer than the budget is passed through unsplit —
|
||||
/// long single sentences are fine (the LM cap bounds runaway); only seams
|
||||
/// are being minimized.
|
||||
#[test]
|
||||
fn chunk_grouping_oversized_sentence_passes_through() {
|
||||
let long = "word ".repeat(60).trim_end().to_string() + ".";
|
||||
assert!(long.len() > 200);
|
||||
let chunks = group_sentences_into_chunks(&s(&["First.", &long]), 200);
|
||||
assert_eq!(chunks, vec!["First.".to_string(), long]);
|
||||
}
|
||||
|
||||
/// Single-sentence messages — the common huddle case, since agents are
|
||||
/// prompted to send one sentence per message — are unaffected by grouping.
|
||||
#[test]
|
||||
fn chunk_grouping_single_sentence_unchanged() {
|
||||
let chunks = group_sentences_into_chunks(&s(&["Just one sentence here."]), 200);
|
||||
assert_eq!(chunks, vec!["Just one sentence here.".to_string()]);
|
||||
}
|
||||
|
||||
/// Empty and whitespace-only entries are dropped, and never produce
|
||||
/// empty chunks (which would synthesize as garbage).
|
||||
#[test]
|
||||
fn chunk_grouping_skips_blank_sentences() {
|
||||
let chunks = group_sentences_into_chunks(&s(&["", " ", "Real sentence.", " ", "Two."]), 200);
|
||||
assert_eq!(chunks[0], "Real sentence.");
|
||||
assert_eq!(chunks.len(), 2);
|
||||
assert_eq!(chunks[1], "Two.");
|
||||
}
|
||||
|
||||
/// Empty input produces no chunks (the worker loop then synthesizes nothing).
|
||||
#[test]
|
||||
fn chunk_grouping_empty_input() {
|
||||
assert!(group_sentences_into_chunks(&[], 200).is_empty());
|
||||
}
|
||||
|
||||
/// Chunks joined with a single space preserve each sentence's terminal
|
||||
/// punctuation — the model sees natural multi-sentence prose, matching the
|
||||
/// shape upstream's ~50-token chunker produces.
|
||||
#[test]
|
||||
fn chunk_grouping_preserves_punctuation_at_joins() {
|
||||
let chunks = group_sentences_into_chunks(&s(&["Lead.", "Really?", "Yes!", "Good."]), 200);
|
||||
assert_eq!(chunks[1], "Really? Yes! Good.");
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user