1#![forbid(unsafe_code)]
4#![warn(missing_docs)]
5
6mod checkpoint;
7mod distribution;
8mod evidence;
9mod parity;
10mod realtime;
11
12pub use checkpoint::{
13 compare_checkpoint_artifacts, CheckpointParityError, CheckpointParityOptions,
14 CheckpointParityReport,
15};
16pub use distribution::{compare_distributions, DistributionError, DistributionMetrics};
17
18pub use evidence::{
19 observe_f32_tensor, observe_i32_tensor, observe_realtime_frame, summarize_latencies,
20 EvaluationEvidence, EvidenceError, LatencySummary,
21};
22pub use parity::{
23 compare_observations, LogitRowMetrics, LogitTolerance, NumericMetrics, NumericTolerance,
24 ObservationParity, ParityComparison, ParityError, ParityMetrics, ParityPolicy, ParityReport,
25 ParityRule,
26};
27pub use realtime::{
28 encoded_audio_frames, run_realtime_trace, RealtimeEvaluationDriver, RealtimeTrace,
29 RealtimeTraceError,
30};
31
32use std::{
33 error::Error,
34 fs,
35 io::{self, Write},
36 path::{Path, PathBuf},
37 time::Instant,
38};
39
40use eredu_architectures::moshi::personaplex_prompt::{
41 materialize_prompt_frames, released_conditioning_plan, wrap_system_prompt,
42 AUDIO_TOKENS_PER_STREAM, RELEASED_PROMPT_SILENCE_FRAMES,
43};
44use eredu_codec::mimi::Mimi;
45use eredu_core::{RealtimeInputFrame, RealtimeOutputFrame, RealtimeSampling};
46use eredu_nn::Tensor;
47use sentencepiece_rs::SentencePieceProcessor;
48use serde::Serialize;
49use serde_json::json;
50
51const SAMPLE_RATE: u32 = 24_000;
52const FRAME_RATE: f64 = 12.5;
53const FRAME_SAMPLES: usize = 1_920;
54const DEADLINE_MS: f64 = 1_000.0 / FRAME_RATE;
55const TAIL_ACTIVITY_FRAMES: usize = 3;
56const ACTIVE_AUDIO_DBFS: f64 = -40.0;
57
58pub const DEFAULT_TEXT_PROMPT: &str = "You are a wise and friendly teacher. Answer questions or provide advice in a clear and engaging way.";
60pub const DEFAULT_SAMPLING_SEED: u64 = 20_260_713;
62
63#[derive(Debug, Clone)]
65pub struct PersonaPlexEvaluationPaths {
66 pub dense_model: PathBuf,
68 pub quantized_model: PathBuf,
70 pub text_tokenizer: PathBuf,
72 pub voice_prompt: PathBuf,
74 pub input: PathBuf,
76 pub output: PathBuf,
78}
79
80#[derive(Debug, Clone)]
82pub struct PersonaPlexEvaluationOptions {
83 pub frames: Option<usize>,
85 pub text_prompt: String,
87 pub sampling_seed: u64,
89}
90
91impl Default for PersonaPlexEvaluationOptions {
92 fn default() -> Self {
93 Self {
94 frames: None,
95 text_prompt: DEFAULT_TEXT_PROMPT.into(),
96 sampling_seed: DEFAULT_SAMPLING_SEED,
97 }
98 }
99}
100
101pub fn run_personaplex_quantization<D, T, L>(
107 paths: &PersonaPlexEvaluationPaths,
108 options: &PersonaPlexEvaluationOptions,
109 mimi: &mut Mimi<T>,
110 context: &T::Context,
111 mut load_model: L,
112) -> Result<(), Box<dyn Error>>
113where
114 D: RealtimeEvaluationDriver,
115 T: Tensor,
116 L: FnMut(&Path) -> Result<D, Box<dyn Error>>,
117{
118 if paths.output.exists() {
119 return Err(invalid(format!(
120 "output directory already exists: {}",
121 paths.output.display()
122 )));
123 }
124 let voice_pcm = read_f32le(&paths.voice_prompt)?;
125 if voice_pcm.len() < FRAME_SAMPLES {
126 return Err(invalid("voice prompt contains no complete 80 ms frame"));
127 }
128 let input_pcm = read_f32le(&paths.input)?;
129 let available_frames = input_pcm.len() / FRAME_SAMPLES;
130 let frames = options
131 .frames
132 .unwrap_or(available_frames)
133 .min(available_frames);
134 if frames < 4 {
135 return Err(invalid(format!(
136 "input must contain at least four complete frames; found {available_frames}"
137 )));
138 }
139 let input_pcm = &input_pcm[..frames * FRAME_SAMPLES];
140 let input_tail = tail_max_rms_dbfs(input_pcm);
141 let input_likely_truncated = input_tail > ACTIVE_AUDIO_DBFS;
142 let input_warning = input_likely_truncated.then_some(
143 "The final 240 ms contains active audio; the frame limit may truncate the user utterance.",
144 );
145 if let Some(warning) = input_warning {
146 eprintln!("warning: {warning} tail_max_rms_dbfs={input_tail:.1}");
147 }
148
149 let codec_start = Instant::now();
150 let voice_tokens = encode_pcm(mimi, &voice_pcm, context)?;
151 let input_tokens = encode_pcm(mimi, input_pcm, context)?;
152 if input_tokens.len() != frames {
153 return Err(invalid(format!(
154 "Mimi produced {} frames for {frames} PCM frames",
155 input_tokens.len()
156 )));
157 }
158 let encode_seconds = codec_start.elapsed().as_secs_f64();
159 let offline = offline_roundtrip(mimi, input_pcm, context)?;
160 let offline_tokens = offline.tokens;
161 let offline_roundtrip = offline.pcm;
162 let streaming_roundtrip = decode_tokens(mimi, &input_tokens, context, input_pcm.len())?;
163 let codec_agreement = token_frame_agreement(&input_tokens, &offline_tokens);
164
165 let tokenizer = SentencePieceProcessor::open(&paths.text_tokenizer)?;
166 let wrapped_text_prompt = wrap_system_prompt(&options.text_prompt);
167 let text_tokens = tokenizer
168 .encode_to_ids(&wrapped_text_prompt)?
169 .into_iter()
170 .map(i32::try_from)
171 .collect::<Result<Vec<_>, _>>()?;
172 if text_tokens.is_empty() {
173 return Err(invalid("text prompt tokenized to an empty sequence"));
174 }
175 let prompt = PromptConditioning {
176 voice_frames: voice_tokens,
177 text_tokens,
178 };
179
180 let dense_load_start = Instant::now();
181 let mut dense = load_model(&paths.dense_model)?;
182 let dense_load_seconds = dense_load_start.elapsed().as_secs_f64();
183 validate_personaplex_geometry(&dense)?;
184 let dense_reference = run_model(
185 &mut dense,
186 &prompt,
187 &input_tokens,
188 RealtimeSampling::greedy(),
189 RunMode::Diagnostics,
190 )?;
191 let sampling =
192 RealtimeSampling::new(0.7, 0.8, options.sampling_seed)?.with_top_k(Some(25), Some(250))?;
193 let dense_run = run_model(&mut dense, &prompt, &input_tokens, sampling, RunMode::Free)?;
194 drop(dense);
195
196 let quantized_load_start = Instant::now();
197 let mut quantized = load_model(&paths.quantized_model)?;
198 let quantized_load_seconds = quantized_load_start.elapsed().as_secs_f64();
199 validate_personaplex_geometry(&quantized)?;
200 let quantized_teacher = run_model(
201 &mut quantized,
202 &prompt,
203 &input_tokens,
204 RealtimeSampling::greedy(),
205 RunMode::TeacherForced(&dense_reference.frames),
206 )?;
207 let quality = quality_summary(&dense_reference.frames, &quantized_teacher.frames)?;
208 let quantized_run = run_model(
209 &mut quantized,
210 &prompt,
211 &input_tokens,
212 sampling,
213 RunMode::Free,
214 )?;
215
216 let decode_start = Instant::now();
217 let dense_pcm = decode_tokens(mimi, &dense_run.emitted_audio, context, input_pcm.len())?;
218 let quantized_pcm =
219 decode_tokens(mimi, &quantized_run.emitted_audio, context, input_pcm.len())?;
220 let decode_seconds = decode_start.elapsed().as_secs_f64();
221 let dense_tail = tail_max_rms_dbfs(&dense_pcm);
222 let quantized_tail = tail_max_rms_dbfs(&quantized_pcm);
223 let swap = options.sampling_seed & 1 == 1;
224 let (sample_a, sample_b, label_a, label_b, tail_a, tail_b) = if swap {
225 (
226 &quantized_pcm,
227 &dense_pcm,
228 "quantized",
229 "dense",
230 quantized_tail,
231 dense_tail,
232 )
233 } else {
234 (
235 &dense_pcm,
236 &quantized_pcm,
237 "dense",
238 "quantized",
239 dense_tail,
240 quantized_tail,
241 )
242 };
243 let truncated_a = tail_a > ACTIVE_AUDIO_DBFS;
244 let truncated_b = tail_b > ACTIVE_AUDIO_DBFS;
245 if truncated_a || truncated_b {
246 eprintln!(
247 "warning: generated speech is active at the output boundary; sample_a_tail_dbfs={tail_a:.1} sample_b_tail_dbfs={tail_b:.1}"
248 );
249 }
250 let dense_performance = performance_summary(&dense_run.latencies_ms);
251 let quantized_performance = performance_summary(&quantized_run.latencies_ms);
252 let divergence = free_run_agreement(&dense_run.frames, &quantized_run.frames);
253
254 fs::create_dir(&paths.output)?;
255 write_wav_pcm16(&paths.output.join("input.wav"), input_pcm, SAMPLE_RATE)?;
256 write_wav_pcm16(
257 &paths.output.join("input_codec_roundtrip.wav"),
258 &streaming_roundtrip,
259 SAMPLE_RATE,
260 )?;
261 write_wav_pcm16(
262 &paths.output.join("input_codec_roundtrip_offline.wav"),
263 &offline_roundtrip,
264 SAMPLE_RATE,
265 )?;
266 write_wav_pcm16(&paths.output.join("sample_a.wav"), sample_a, SAMPLE_RATE)?;
267 write_wav_pcm16(&paths.output.join("sample_b.wav"), sample_b, SAMPLE_RATE)?;
268
269 let metrics = json!({
270 "format_version": 1,
271 "methodology": "Both models use the portable realtime evaluation driver, forcing, sampling, and observation contracts.",
272 "input": {
273 "path": paths.input,
274 "sample_rate": SAMPLE_RATE,
275 "frame_rate": FRAME_RATE,
276 "frames": frames,
277 "audio_seconds": frames as f64 / FRAME_RATE,
278 "tail_max_rms_dbfs": input_tail,
279 "likely_truncated": input_likely_truncated,
280 "warning": input_warning,
281 },
282 "conditioning": {
283 "voice_prompt_path": paths.voice_prompt,
284 "voice_prompt_frames": prompt.voice_frames.len(),
285 "text_tokenizer_path": paths.text_tokenizer,
286 "text_prompt": options.text_prompt,
287 "wrapped_text_prompt": wrapped_text_prompt,
288 "text_prompt_tokens": prompt.text_tokens.len(),
289 "silence_frames_after_voice": RELEASED_PROMPT_SILENCE_FRAMES,
290 "silence_frames_after_text": RELEASED_PROMPT_SILENCE_FRAMES,
291 },
292 "codec_diagnostic": {
293 "streaming_roundtrip": "input_codec_roundtrip.wav",
294 "offline_roundtrip": "input_codec_roundtrip_offline.wav",
295 "streaming_offline_token_agreement": codec_agreement,
296 },
297 "performance": {
298 "frame_deadline_ms": DEADLINE_MS,
299 "codec_encode_seconds": encode_seconds,
300 "codec_decode_both_outputs_seconds": decode_seconds,
301 "dense": { "load_seconds": dense_load_seconds, "model": dense_performance },
302 "quantized": { "load_seconds": quantized_load_seconds, "model": quantized_performance },
303 },
304 "teacher_forced_quality": quality,
305 "free_run_divergence_diagnostic": divergence,
306 "listening_test": {
307 "input": "input.wav",
308 "sample_a": "sample_a.wav",
309 "sample_b": "sample_b.wav",
310 "sampling": {
311 "seed": options.sampling_seed,
312 "text_temperature": sampling.text_temperature(),
313 "audio_temperature": sampling.audio_temperature(),
314 "text_top_k": sampling.text_top_k(),
315 "audio_top_k": sampling.audio_top_k(),
316 },
317 "sample_a_tail_max_rms_dbfs": tail_a,
318 "sample_b_tail_max_rms_dbfs": tail_b,
319 "sample_a_likely_truncated": truncated_a,
320 "sample_b_likely_truncated": truncated_b,
321 "input_warning": input_warning,
322 },
323 });
324 fs::write(
325 paths.output.join("metrics.json"),
326 serde_json::to_vec_pretty(&metrics)?,
327 )?;
328 fs::write(
329 paths.output.join("answer_key.json"),
330 serde_json::to_vec_pretty(&json!({ "sample_a": label_a, "sample_b": label_b }))?,
331 )?;
332 fs::write(
333 paths.output.join("listening_manifest.json"),
334 serde_json::to_vec_pretty(&json!({
335 "format_version": 1,
336 "trials": [{
337 "id": "personaplex_quantization_001",
338 "input": "input.wav",
339 "codec_roundtrip": "input_codec_roundtrip.wav",
340 "sample_a": "sample_a.wav",
341 "sample_b": "sample_b.wav",
342 "input_warning": input_warning,
343 "sample_a_likely_truncated": truncated_a,
344 "sample_b_likely_truncated": truncated_b,
345 }],
346 }))?,
347 )?;
348 fs::write(
349 paths.output.join("token_diagnostics.json"),
350 serde_json::to_vec_pretty(&json!({
351 "input": input_tokens,
352 "input_offline": offline_tokens,
353 "conditioning": {
354 "voice_prompt": prompt.voice_frames,
355 "text_prompt": prompt.text_tokens,
356 "silence_frames_after_voice": RELEASED_PROMPT_SILENCE_FRAMES,
357 "silence_frames_after_text": RELEASED_PROMPT_SILENCE_FRAMES,
358 },
359 "sampling": {
360 "seed": options.sampling_seed,
361 "text_temperature": sampling.text_temperature(),
362 "audio_temperature": sampling.audio_temperature(),
363 "text_top_k": sampling.text_top_k(),
364 "audio_top_k": sampling.audio_top_k(),
365 },
366 "dense_emitted": dense_run.emitted_audio,
367 "dense_sampled_frames": reference_tokens(&dense_run.frames),
368 "dense_greedy_emitted": dense_reference.emitted_audio,
369 "dense_greedy_frames": reference_tokens(&dense_reference.frames),
370 "quantized_emitted": quantized_run.emitted_audio,
371 }))?,
372 )?;
373 Ok(())
374}
375
376fn validate_personaplex_geometry<D: RealtimeEvaluationDriver>(
377 driver: &D,
378) -> Result<(), Box<dyn Error>> {
379 let config = driver.speech_config();
380 if config.input_audio_codebooks() != AUDIO_TOKENS_PER_STREAM
381 || config.generated_audio_codebooks() != AUDIO_TOKENS_PER_STREAM
382 {
383 return Err(invalid(format!(
384 "PersonaPlex evaluation requires {AUDIO_TOKENS_PER_STREAM} input and generated codebooks, got {} and {}",
385 config.input_audio_codebooks(),
386 config.generated_audio_codebooks()
387 )));
388 }
389 Ok(())
390}
391
392struct PromptConditioning {
393 voice_frames: Vec<Vec<i32>>,
394 text_tokens: Vec<i32>,
395}
396
397enum RunMode<'a> {
398 Free,
399 Diagnostics,
400 TeacherForced(&'a [ReferenceFrame]),
401}
402
403struct ReferenceFrame {
404 text_token: i32,
405 decision_audio: Vec<i32>,
406 sampled_audio: Vec<i32>,
407 diagnostics: Vec<Vec<f32>>,
408}
409
410struct ModelRun {
411 frames: Vec<ReferenceFrame>,
412 emitted_audio: Vec<Vec<i32>>,
413 latencies_ms: Vec<f64>,
414}
415
416fn run_model<D: RealtimeEvaluationDriver>(
417 driver: &mut D,
418 prompt: &PromptConditioning,
419 input_tokens: &[Vec<i32>],
420 sampling: RealtimeSampling,
421 mode: RunMode<'_>,
422) -> Result<ModelRun, Box<dyn Error>> {
423 driver
424 .start_trace(sampling)
425 .map_err(evaluation_driver_error::<D::Error>)?;
426 for frame in prompt_frames(prompt)? {
427 run_frame(driver, frame)?;
428 }
429 let mut frames = Vec::with_capacity(input_tokens.len());
430 let mut emitted_audio = Vec::new();
431 let mut latencies_ms = Vec::with_capacity(input_tokens.len());
432 for (index, tokens) in input_tokens.iter().enumerate() {
433 let mut frame = RealtimeInputFrame::new(1, tokens.clone());
434 match mode {
435 RunMode::Free => {}
436 RunMode::Diagnostics => frame = frame.with_diagnostics(),
437 RunMode::TeacherForced(reference) => {
438 let reference = reference
439 .get(index)
440 .ok_or_else(|| invalid("teacher-forced reference is shorter than input"))?;
441 frame = frame
442 .with_forced_text(vec![reference.text_token])
443 .with_forced_generated_audio(reference.sampled_audio.clone())
444 .with_diagnostics();
445 }
446 }
447 let start = Instant::now();
448 let output = run_frame(driver, frame)?;
449 latencies_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
450 if let Some(tokens) = output.output_audio_tokens() {
451 emitted_audio.push(tokens.to_vec());
452 }
453 frames.push(ReferenceFrame {
454 text_token: *output
455 .text_tokens()
456 .first()
457 .ok_or_else(|| invalid("realtime output has no text token"))?,
458 decision_audio: output.decision_audio_tokens().to_vec(),
459 sampled_audio: output.sampled_audio_tokens().to_vec(),
460 diagnostics: output
461 .diagnostics()
462 .iter()
463 .map(|diagnostic| diagnostic.logits().to_vec())
464 .collect(),
465 });
466 }
467 driver
468 .finish_trace()
469 .map_err(evaluation_driver_error::<D::Error>)?;
470 Ok(ModelRun {
471 frames,
472 emitted_audio,
473 latencies_ms,
474 })
475}
476
477fn run_frame<D: RealtimeEvaluationDriver>(
478 driver: &mut D,
479 frame: RealtimeInputFrame,
480) -> Result<RealtimeOutputFrame, Box<dyn Error>> {
481 driver
482 .evaluate_frame(frame)
483 .map_err(evaluation_driver_error::<D::Error>)
484}
485
486fn evaluation_driver_error<E>(error: E) -> Box<dyn Error>
487where
488 E: Error + Send + Sync + 'static,
489{
490 Box::new(error)
491}
492
493fn prompt_frames(prompt: &PromptConditioning) -> Result<Vec<RealtimeInputFrame>, Box<dyn Error>> {
494 if let Some(frame) = prompt
495 .voice_frames
496 .iter()
497 .find(|frame| frame.len() != AUDIO_TOKENS_PER_STREAM)
498 {
499 return Err(invalid(format!(
500 "PersonaPlex voice prompt frame must contain {AUDIO_TOKENS_PER_STREAM} tokens, got {}",
501 frame.len()
502 )));
503 }
504 let voice_frame_count = prompt.voice_frames.len();
505 let mut voice = Vec::with_capacity(AUDIO_TOKENS_PER_STREAM * voice_frame_count);
506 for codebook in 0..AUDIO_TOKENS_PER_STREAM {
507 voice.extend(prompt.voice_frames.iter().map(|frame| frame[codebook]));
508 }
509 let voice_shape =
510 (voice_frame_count > 0).then_some([1, AUDIO_TOKENS_PER_STREAM, voice_frame_count]);
511 let plan = released_conditioning_plan(
512 voice_shape.as_ref().map(|shape| shape.as_slice()),
513 &[1, prompt.text_tokens.len()],
514 )?;
515 Ok(materialize_prompt_frames(
516 &plan,
517 &voice,
518 voice_frame_count,
519 &prompt.text_tokens,
520 prompt.text_tokens.len(),
521 )?)
522}
523
524fn encode_pcm<T: Tensor>(
525 mimi: &mut Mimi<T>,
526 pcm: &[f32],
527 context: &T::Context,
528) -> Result<Vec<Vec<i32>>, Box<dyn Error>> {
529 mimi.reset_encode_state();
530 let mut frames = Vec::with_capacity(pcm.len() / FRAME_SAMPLES);
531 for frame in pcm.as_chunks::<FRAME_SAMPLES>().0 {
532 let frame = T::from_f32_slice(frame, &[1, 1, FRAME_SAMPLES as i32], context)?;
533 if let Some(tokens) = mimi.encode_step(&frame, context)? {
534 frames.push(tokens.to_i32_vec(context)?);
535 }
536 }
537 Ok(frames)
538}
539
540fn decode_tokens<T: Tensor>(
541 mimi: &mut Mimi<T>,
542 frames: &[Vec<i32>],
543 context: &T::Context,
544 target_samples: usize,
545) -> Result<Vec<f32>, Box<dyn Error>> {
546 mimi.reset_decode_state();
547 let mut pcm = Vec::with_capacity(target_samples);
548 for frame in frames {
549 let tokens = T::from_i32_slice(frame, &[1, frame.len() as i32], context)?;
550 pcm.extend(mimi.decode_step(&tokens, context)?.to_f32_vec(context)?);
551 }
552 pcm.truncate(target_samples);
553 pcm.resize(target_samples, 0.0);
554 Ok(pcm)
555}
556
557fn offline_roundtrip<T: Tensor>(
558 mimi: &mut Mimi<T>,
559 pcm: &[f32],
560 context: &T::Context,
561) -> Result<OfflineRoundtrip, Box<dyn Error>> {
562 let input = T::from_f32_slice(pcm, &[1, 1, pcm.len() as i32], context)?;
563 let codes = mimi.encode(&input, context)?;
564 let code_shape = codes.shape().to_vec();
565 if code_shape.len() != 3 || code_shape[0] != 1 {
566 return Err(invalid(format!(
567 "offline Mimi codes have unexpected shape {code_shape:?}"
568 )));
569 }
570 let values = codes.to_i32_vec(context)?;
571 let codebooks = code_shape[1] as usize;
572 let frame_count = code_shape[2] as usize;
573 let mut frames = vec![vec![0; codebooks]; frame_count];
574 for codebook in 0..codebooks {
575 for frame in 0..frame_count {
576 frames[frame][codebook] = values[codebook * frame_count + frame];
577 }
578 }
579 let mut roundtrip = mimi.decode(&codes, context)?.to_f32_vec(context)?;
580 roundtrip.truncate(pcm.len());
581 roundtrip.resize(pcm.len(), 0.0);
582 Ok(OfflineRoundtrip {
583 tokens: frames,
584 pcm: roundtrip,
585 })
586}
587
588struct OfflineRoundtrip {
589 tokens: Vec<Vec<i32>>,
590 pcm: Vec<f32>,
591}
592
593#[derive(Debug, Clone, Default)]
594struct DistributionAccumulator {
595 count: usize,
596 target_count: usize,
597 kl_sum: f64,
598 entropy_sum: f64,
599 target_nll_delta_sum: f64,
600 centered_rmse_sum: f64,
601 top1_matches: usize,
602 top5_overlap_sum: f64,
603}
604
605impl DistributionAccumulator {
606 fn update(
607 &mut self,
608 dense: &[f32],
609 candidate: &[f32],
610 target: usize,
611 ) -> Result<(), Box<dyn Error>> {
612 let metrics = compare_distributions(
613 dense,
614 candidate,
615 (target < dense.len()).then_some(target),
616 5,
617 )?;
618 self.count += 1;
619 self.kl_sum += metrics.kl_nats;
620 self.entropy_sum += metrics.reference_entropy_nats;
621 self.centered_rmse_sum += metrics.centered_logit_rmse;
622 self.top1_matches += usize::from(metrics.top1_agreement);
623 self.top5_overlap_sum += metrics.top_k_overlap;
624 if let Some(delta) = metrics.target_nll_delta_nats {
625 self.target_count += 1;
626 self.target_nll_delta_sum += delta;
627 }
628 Ok(())
629 }
630
631 fn merge(&mut self, other: &Self) {
632 self.count += other.count;
633 self.target_count += other.target_count;
634 self.kl_sum += other.kl_sum;
635 self.entropy_sum += other.entropy_sum;
636 self.target_nll_delta_sum += other.target_nll_delta_sum;
637 self.centered_rmse_sum += other.centered_rmse_sum;
638 self.top1_matches += other.top1_matches;
639 self.top5_overlap_sum += other.top5_overlap_sum;
640 }
641
642 fn summary(&self) -> MetricSummary {
643 let count = self.count.max(1) as f64;
644 MetricSummary {
645 distributions: self.count,
646 target_distributions: self.target_count,
647 mean_kl_nats: self.kl_sum / count,
648 mean_dense_entropy_nats: self.entropy_sum / count,
649 mean_target_nll_delta_nats: self.target_nll_delta_sum / self.target_count.max(1) as f64,
650 mean_centered_logit_rmse: self.centered_rmse_sum / count,
651 top1_agreement: self.top1_matches as f64 / count,
652 mean_top5_overlap: self.top5_overlap_sum / count,
653 }
654 }
655}
656
657#[derive(Debug, Clone, Serialize)]
658struct MetricSummary {
659 distributions: usize,
660 target_distributions: usize,
661 mean_kl_nats: f64,
662 mean_dense_entropy_nats: f64,
663 mean_target_nll_delta_nats: f64,
664 mean_centered_logit_rmse: f64,
665 top1_agreement: f64,
666 mean_top5_overlap: f64,
667}
668
669#[derive(Debug, Clone, Serialize)]
670struct QualitySummary {
671 methodology: &'static str,
672 text: MetricSummary,
673 audio_generated: MetricSummary,
674 audio_input_conditioned: MetricSummary,
675 audio_overall: MetricSummary,
676 audio_by_codebook: Vec<MetricSummary>,
677}
678
679fn quality_summary(
680 dense: &[ReferenceFrame],
681 candidate: &[ReferenceFrame],
682) -> Result<QualitySummary, Box<dyn Error>> {
683 if dense.len() != candidate.len() {
684 return Err(invalid("teacher-forced run lengths differ"));
685 }
686 let mut text = DistributionAccumulator::default();
687 let mut audio = Vec::<DistributionAccumulator>::new();
688 for (dense, candidate) in dense.iter().zip(candidate) {
689 if dense.diagnostics.len() != candidate.diagnostics.len() || dense.diagnostics.is_empty() {
690 return Err(invalid(
691 "teacher-forced diagnostic counts differ or are empty",
692 ));
693 }
694 text.update(
695 &dense.diagnostics[0],
696 &candidate.diagnostics[0],
697 dense.text_token as usize,
698 )?;
699 if audio.is_empty() {
700 audio.resize(
701 dense.diagnostics.len() - 1,
702 DistributionAccumulator::default(),
703 );
704 }
705 for (codebook, accumulator) in audio.iter_mut().enumerate() {
706 accumulator.update(
707 &dense.diagnostics[codebook + 1],
708 &candidate.diagnostics[codebook + 1],
709 *dense
710 .decision_audio
711 .get(codebook)
712 .ok_or_else(|| invalid("teacher-forced decision token is missing"))?
713 as usize,
714 )?;
715 }
716 }
717 let mut overall = DistributionAccumulator::default();
718 for value in &audio {
719 overall.merge(value);
720 }
721 let mut generated = DistributionAccumulator::default();
722 for value in audio.iter().take(AUDIO_TOKENS_PER_STREAM) {
723 generated.merge(value);
724 }
725 let mut input_conditioned = DistributionAccumulator::default();
726 for value in audio.iter().skip(AUDIO_TOKENS_PER_STREAM) {
727 input_conditioned.merge(value);
728 }
729 Ok(QualitySummary {
730 methodology: "The candidate is teacher-forced onto the dense model's exact text and generated-audio history; KL uses the dense distribution as reference.",
731 text: text.summary(),
732 audio_generated: generated.summary(),
733 audio_input_conditioned: input_conditioned.summary(),
734 audio_overall: overall.summary(),
735 audio_by_codebook: audio.iter().map(DistributionAccumulator::summary).collect(),
736 })
737}
738
739#[derive(Debug, Clone, Serialize)]
740struct PerformanceSummary {
741 frames: usize,
742 mean_ms: f64,
743 p50_ms: f64,
744 p95_ms: f64,
745 max_ms: f64,
746 deadline_misses: usize,
747}
748
749fn performance_summary(latencies: &[f64]) -> PerformanceSummary {
750 let summary = summarize_latencies(latencies, Some(DEADLINE_MS))
751 .expect("every model run records at least one finite nonnegative latency");
752 PerformanceSummary {
753 frames: summary.samples,
754 mean_ms: summary.mean_ms,
755 p50_ms: summary.p50_ms,
756 p95_ms: summary.p95_ms,
757 max_ms: summary.max_ms,
758 deadline_misses: summary.deadline_misses,
759 }
760}
761
762fn free_run_agreement(dense: &[ReferenceFrame], quantized: &[ReferenceFrame]) -> serde_json::Value {
763 let frames = dense.len().min(quantized.len());
764 let text_matches = dense
765 .iter()
766 .zip(quantized)
767 .filter(|(left, right)| left.text_token == right.text_token)
768 .count();
769 let mut audio_matches = 0usize;
770 let mut audio_total = 0usize;
771 for (left, right) in dense.iter().zip(quantized) {
772 for (left, right) in left.sampled_audio.iter().zip(&right.sampled_audio) {
773 audio_matches += usize::from(left == right);
774 audio_total += 1;
775 }
776 }
777 json!({
778 "frames": frames,
779 "text_token_agreement": text_matches as f64 / frames.max(1) as f64,
780 "audio_token_agreement": audio_matches as f64 / audio_total.max(1) as f64,
781 })
782}
783
784fn reference_tokens(frames: &[ReferenceFrame]) -> Vec<serde_json::Value> {
785 frames
786 .iter()
787 .map(|frame| json!({ "text": frame.text_token, "sampled_audio": frame.sampled_audio }))
788 .collect()
789}
790
791fn token_frame_agreement(left: &[Vec<i32>], right: &[Vec<i32>]) -> f64 {
792 let mut matches = 0usize;
793 let mut total = 0usize;
794 for (left, right) in left.iter().zip(right) {
795 for (left, right) in left.iter().zip(right) {
796 matches += usize::from(left == right);
797 total += 1;
798 }
799 }
800 matches as f64 / total.max(1) as f64
801}
802
803fn read_f32le(path: &Path) -> Result<Vec<f32>, Box<dyn Error>> {
804 let bytes = fs::read(path)?;
805 if bytes.len() % 4 != 0 {
806 return Err(invalid(format!(
807 "raw f32le input length must be divisible by four, got {} bytes",
808 bytes.len()
809 )));
810 }
811 Ok(bytes
812 .as_chunks::<4>()
813 .0
814 .iter()
815 .map(|chunk| f32::from_le_bytes(*chunk))
816 .collect())
817}
818
819fn rms_dbfs(samples: &[f32]) -> f64 {
820 let mean_square = samples
821 .iter()
822 .map(|sample| (*sample as f64) * (*sample as f64))
823 .sum::<f64>()
824 / samples.len().max(1) as f64;
825 20.0 * mean_square.sqrt().max(1e-12).log10()
826}
827
828fn tail_max_rms_dbfs(samples: &[f32]) -> f64 {
829 samples
830 .as_chunks::<FRAME_SAMPLES>()
831 .0
832 .iter()
833 .rev()
834 .take(TAIL_ACTIVITY_FRAMES)
835 .map(|frame| rms_dbfs(frame))
836 .fold(f64::NEG_INFINITY, f64::max)
837}
838
839fn write_wav_pcm16(path: &Path, samples: &[f32], sample_rate: u32) -> Result<(), Box<dyn Error>> {
840 let data_bytes = u32::try_from(
841 samples
842 .len()
843 .checked_mul(2)
844 .ok_or_else(|| invalid("WAV size overflow"))?,
845 )?;
846 let mut file = fs::File::create(path)?;
847 file.write_all(b"RIFF")?;
848 file.write_all(&(36u32 + data_bytes).to_le_bytes())?;
849 file.write_all(b"WAVEfmt ")?;
850 file.write_all(&16u32.to_le_bytes())?;
851 file.write_all(&1u16.to_le_bytes())?;
852 file.write_all(&1u16.to_le_bytes())?;
853 file.write_all(&sample_rate.to_le_bytes())?;
854 file.write_all(&(sample_rate * 2).to_le_bytes())?;
855 file.write_all(&2u16.to_le_bytes())?;
856 file.write_all(&16u16.to_le_bytes())?;
857 file.write_all(b"data")?;
858 file.write_all(&data_bytes.to_le_bytes())?;
859 for sample in samples {
860 let value = (sample.clamp(-1.0, 1.0) * i16::MAX as f32).round() as i16;
861 file.write_all(&value.to_le_bytes())?;
862 }
863 Ok(())
864}
865
866fn invalid(message: impl Into<String>) -> Box<dyn Error> {
867 Box::new(io::Error::new(io::ErrorKind::InvalidInput, message.into()))
868}
869
870#[cfg(test)]
871mod tests {
872 use super::*;
873 use eredu_core::{RealtimeFrameConvention, RealtimeSpeechConfig};
874
875 struct PersonaRecordingDriver {
876 config: RealtimeSpeechConfig,
877 sampling: Vec<RealtimeSampling>,
878 frames: Vec<RealtimeInputFrame>,
879 finishes: usize,
880 }
881
882 impl PersonaRecordingDriver {
883 fn new() -> Self {
884 Self {
885 config: RealtimeSpeechConfig::new(
886 16,
887 8,
888 8,
889 8,
890 eredu_architectures::moshi::personaplex_prompt::TEXT_PADDING_TOKEN,
891 2_048,
892 RealtimeFrameConvention::FeedbackAlignedHistory,
893 vec![0; 17],
894 )
895 .unwrap(),
896 sampling: Vec::new(),
897 frames: Vec::new(),
898 finishes: 0,
899 }
900 }
901 }
902
903 impl RealtimeEvaluationDriver for PersonaRecordingDriver {
904 type Error = std::io::Error;
905
906 fn speech_config(&self) -> &RealtimeSpeechConfig {
907 &self.config
908 }
909
910 fn start_trace(&mut self, sampling: RealtimeSampling) -> Result<(), Self::Error> {
911 self.sampling.push(sampling);
912 Ok(())
913 }
914
915 fn evaluate_frame(
916 &mut self,
917 frame: RealtimeInputFrame,
918 ) -> Result<RealtimeOutputFrame, Self::Error> {
919 let text = frame
920 .forced_text_tokens()
921 .map_or_else(|| vec![7], <[i32]>::to_vec);
922 let audio = frame
923 .forced_generated_audio_tokens()
924 .map_or_else(|| vec![9; 8], <[i32]>::to_vec);
925 self.frames.push(frame);
926 Ok(RealtimeOutputFrame::new(
927 1,
928 text,
929 audio.clone(),
930 audio.clone(),
931 Some(audio),
932 Vec::new(),
933 ))
934 }
935
936 fn finish_trace(&mut self) -> Result<(), Self::Error> {
937 self.finishes += 1;
938 Ok(())
939 }
940 }
941
942 #[test]
943 fn identical_distribution_metrics_are_exact() {
944 let values = [0.0, 1.0, -1.0, 0.5, 0.25];
945 let mut metric = DistributionAccumulator::default();
946 metric.update(&values, &values, 1).unwrap();
947 let summary = metric.summary();
948 assert!(summary.mean_kl_nats.abs() < 1e-12);
949 assert_eq!(summary.top1_agreement, 1.0);
950 assert_eq!(summary.mean_top5_overlap, 1.0);
951 }
952
953 #[test]
954 fn prompt_frames_preserve_released_conditioning_order() {
955 let prompt = PromptConditioning {
956 voice_frames: vec![vec![1; AUDIO_TOKENS_PER_STREAM]],
957 text_tokens: vec![7, 8],
958 };
959 let frames = prompt_frames(&prompt).unwrap();
960 assert_eq!(
961 frames.len(),
962 1 + RELEASED_PROMPT_SILENCE_FRAMES + 2 + RELEASED_PROMPT_SILENCE_FRAMES
963 );
964 assert_eq!(frames[0].forced_generated_audio_tokens(), Some(&[1; 8][..]));
965 assert_eq!(
966 frames[1 + RELEASED_PROMPT_SILENCE_FRAMES].forced_text_tokens(),
967 Some(&[7][..])
968 );
969 }
970
971 #[test]
972 fn personaplex_run_uses_the_evaluation_driver_for_prompt_and_live_frames() {
973 let prompt = PromptConditioning {
974 voice_frames: vec![vec![1; AUDIO_TOKENS_PER_STREAM]],
975 text_tokens: vec![7, 8],
976 };
977 let expected_prompt_frames = prompt_frames(&prompt).unwrap().len();
978 let mut driver = PersonaRecordingDriver::new();
979
980 let run = run_model(
981 &mut driver,
982 &prompt,
983 &[vec![3; AUDIO_TOKENS_PER_STREAM]],
984 RealtimeSampling::greedy(),
985 RunMode::Free,
986 )
987 .unwrap();
988
989 assert_eq!(driver.sampling, [RealtimeSampling::greedy()]);
990 assert_eq!(driver.frames.len(), expected_prompt_frames + 1);
991 assert_eq!(driver.finishes, 1);
992 assert_eq!(run.frames.len(), 1);
993 assert_eq!(run.frames[0].text_token, 7);
994 assert_eq!(run.emitted_audio, [vec![9; AUDIO_TOKENS_PER_STREAM]]);
995 }
996}