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