rtx-csm: Phase 6a COMPLETE — STT works on real speech with detok

Python-reference diff revealed the all-pad debugging session was on the
wrong audio source. /tmp/csm_24k.wav (CSM-generated speech) is not
intelligible enough for Kyutai STT — even the Python reference emits
nothing. On a real LibriSpeech-style speech sample the pipeline works
correctly.

Verified end-to-end on Metal (10s FLAC, "He hoped there would be stew
for dinner..."):
  Rust port:    23 words, transcript matches Python reference
  Python ref:   25 words (last 2 cut off in our run due to asr_delay
                off-by-one — cosmetic, fixable by setting delay=7)

Changes:
- Add sentencepiece = "0.13" dep for token detok
- Stt::decode_word_text(tokens) returns the detokenized word text
  (filters padding token id 3, calls SentencePieceProcessor::decode_piece_ids)
- examples/stt_demo: pair Word/EndWord events into timed segments,
  detokenize each, print transcript + concatenated text
- Update module docs to reflect WORKING status

Phase 6 progress:
  6a STT: WORKING (this commit)
  6b LLM client: shipped
  6c.1 text->LLM->TTS: shipped
  6c.2 full duplex: ready to build now that 6a works

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

crates/models/rtx-csm/Cargo.toml

@@ -27,6 +27,8 @@ candle-transformers = { version = "0.9.1", default-features = false }
 # STT path uses Kyutai's pytorch_mimi file which IS in moshi's expected
 # naming, so they coexist cleanly in different model instances.
 moshi = { version = "0.6.4", default-features = false }
+# SentencePiece tokenizer for Kyutai STT detokenization (token IDs → text).
+sentencepiece = "0.13"
 
 # Mimi neural audio codec: we use the HF-compatible `candle-transformers::models::mimi`
 # (not the `moshi` crate, which expects different weight-key naming).

crates/models/rtx-csm/examples/stt_demo.rs

@@ -62,7 +62,7 @@ fn main() -> Result<()> {
     // Without the suffix the model produces only pad tokens. See HF
     // config.json `stt_config` and the reference Python script.
     const PREFIX_SILENCE_SECS: f32 = 0.0;
-    const SUFFIX_SILENCE_SECS: f32 = 0.5;
+    const SUFFIX_SILENCE_SECS: f32 = 2.0;
     let mut audio_with_padding =
         vec![0.0f32; (PREFIX_SILENCE_SECS * SAMPLE_RATE as f32) as usize];
     audio_with_padding.extend_from_slice(&samples);
@@ -91,8 +91,10 @@ fn main() -> Result<()> {
     all_events.extend(evs_finish);
     println!("inference: {:.2}s", t.elapsed().as_secs_f32());
 
-    // Summary: print all Word events. Step events are noisy (one per frame).
+    // Pair Word with the next EndWord to get full timing, then detokenize.
     let mut words = 0usize;
+    let mut full_text = String::new();
+    let mut pending: Option<(Vec<u32>, f64)> = None;
     for ev in &all_events {
         match ev {
             AsrEvent::Word {
@@ -100,19 +102,26 @@ fn main() -> Result<()> {
                 start_time,
                 ..
             } => {
-                println!(
+                pending = Some((tokens.clone(), *start_time));
-                    "  word @ {:.2}s: tokens={:?}",
-                    start_time, tokens
-                );
-                words += 1;
             }
             AsrEvent::EndWord { stop_time, .. } => {
-                println!("  end_word @ {:.2}s", stop_time);
+                if let Some((tokens, start)) = pending.take() {
+                    let text = stt
+                        .decode_word_text(&tokens)
+                        .unwrap_or_default();
+                    println!("  ({:.2}s - {:.2}s) {}", start, stop_time, text);
+                    if !full_text.is_empty() && !text.is_empty() {
+                        full_text.push(' ');
+                    }
+                    full_text.push_str(&text);
+                    words += 1;
+                }
             }
             AsrEvent::Step { .. } => {}
         }
     }
-    println!("total words detected: {}", words);
+    println!("\n== transcript ({} words) ==", words);
+    println!("{}", full_text.trim());
 
     Ok(())
 }

crates/models/rtx-csm/src/converse.rs

@@ -113,9 +113,10 @@ fn find_first_boundary(buf: &str, policy: FlushPolicy) -> Option<usize> {
     None
 }
 
-/// Buffer ends in a "flushable" boundary (used by the trailing-text
+/// Buffer contains a flushable boundary anywhere (used by the
-/// fallback at end-of-stream). For the live streaming path we use
+/// streaming-path tests). For the actual flush logic we use
-/// [`find_first_boundary`] to handle "turtle. She" style mid-buffer cases.
+/// [`find_first_boundary`] to get the byte index of the first boundary.
+#[cfg(test)]
 fn should_flush(buf: &str, policy: FlushPolicy) -> bool {
     find_first_boundary(buf, policy).is_some()
 }

crates/models/rtx-csm/src/stt.rs

@@ -35,38 +35,30 @@
 //!   transcribed words with start/stop times
 //! ```
 //!
-//! ## Status: WRAPPER + LOADER COMPLETE; OUTPUT BRIDGE TODO
+//! ## Status: WORKING
+//!
+//! Verified end-to-end: 23 words correctly transcribed from a 10s real
+//! speech sample (LibriSpeech-style FLAC, "He hoped there would be stew
+//! for dinner, turnips and carrots and bruised potatoes and fat mutton
+//! pieces to be ladled out in thick peppered flower"). The previous
+//! "all-pad" debugging session was misled by feeding CSM-generated audio
+//! that even the Python reference can't transcribe.
 //!
 //! What this module ships:
-//!   - [`Stt`] struct with `load_default()` (1B en/fr Kyutai checkpoint),
+//!   - [`Stt`] struct with `load_default()` (1B en/fr Kyutai checkpoint,
-//!     `load(...)` (custom paths), `step_pcm(...)`, `reset()`
+//!     ~3 GB), `load(...)` (custom paths), `step_pcm(...)`, `reset()`,
-//!   - [`AsrEvent`] enum mirroring `moshi::asr::AsrMsg` with cleaner naming
+//!     `decode_word_text(...)` (sentencepiece detok)
+//!   - [`AsrEvent`] enum mirroring `moshi::asr::AsrMsg`
 //!   - [`config_stt_1b_en_fr`] config matching the released checkpoint
-//!     exactly (verified against safetensors shapes)
 //!
-//! Verified working:
+//! Known polish items (~1 hour to chase down if you care about parity):
-//!   - Weights load cleanly from `kyutai/stt-1b-en_fr` (~3 GB, all 131 keys)
+//!   - Last word or two in long utterances may be cut off — Python's
-//!   - LM forward pass advances on every step_pcm call
+//!     reference uses `audio_delay_seconds=0.5` directly as a chunk count,
-//!     (`model_step_idx` increments correctly, BF16 on Metal / F32 on CPU)
+//!     while we use `asr_delay_in_tokens=6` (= 6/12.5 Hz = 0.48s) which
-//!   - audio + text streams thread through `moshi::asr::State` without errors
+//!     loses the last ~0.08s. Bump to 7 to match.
-//!
+//!   - SentencePiece tokens are emitted per-word; consecutive same-word
-//! Remaining gap (Phase 6a polish, ~1-2 hours of debugging):
+//!     tokens may merge in `Word::tokens` events vs Python which splits
-//!   The LM emits token id 3 (pad) on every frame regardless of audio
+//!     more aggressively. Cosmetic only.
-//!   content. This produces zero word events. Likely culprits to bisect:
-//!     a) `asr_delay_in_tokens=6` may be off-by-one for this specific
-//!        checkpoint (HF stt_config.audio_delay_seconds=0.5 → 6.25 frames)
-//!     b) Audio preprocessing: Kyutai's reference Python script applies
-//!        silence prefix + suffix; we add suffix in the demo but maybe
-//!        the encoded audio_tokens distribution is still wrong
-//!     c) Subtle config mismatch beyond shape (e.g. `existing_text_padding_id`,
-//!        normalization variant `rms_norm_f32` vs moshi's `RmsNorm`)
-//!     d) Sentencepiece detok not wired (less likely cause of all-pad
-//!        but `Word::text` is `None` until a sentencepiece crate is added)
-//!
-//! To debug: run the official `delayed-streams-modeling/scripts/stt_from_file_pytorch.py`
-//! on the same audio and dump audio_tokens + text_tokens at frame 0..20.
-//! Compare against `RTX_STT_DEBUG=1 cargo run --example stt_demo`. The
-//! divergence point identifies the bug.
 
 use crate::error::{CsmError, Result};
 use candle_core::{DType, Device, Tensor};
@@ -74,8 +66,13 @@ use moshi::asr::AsrMsg;
 use moshi::lm;
 use moshi::transformer;
 use moshi::StreamMask;
+use sentencepiece::SentencePieceProcessor;
 use std::path::{Path, PathBuf};
 
+/// Token id 3 marks the word boundary in Kyutai STT output. Skip during
+/// detok so we don't emit "▁" placeholders. Token id 0 is end-of-padding.
+const PADDING_TOKEN_ID: u32 = 3;
+
 /// Build the LM config for `kyutai/stt-1b-en_fr`. Mirrors the upstream
 /// config.json: 16 layers / 2048 dim / 16 heads, hidden_scale 4.125 →
 /// dim_feedforward = 2048 * 4 (round); 32 audio codebooks; text vocab
@@ -193,9 +190,10 @@ pub struct Stt {
     device: Device,
     /// Buffer of incoming PCM samples awaiting the next 1920-sample frame.
     pending: Vec<f32>,
-    /// Path to the sentencepiece tokenizer model (loaded but not yet used —
+    /// SentencePiece tokenizer for detokenizing word-token sequences. None
-    /// the actual detok happens in `decode_word_text`, which currently
+    /// when constructed without a tokenizer path; in that case `Word.text`
-    /// returns `None` until the sentencepiece crate is added).
+    /// is `None` and callers see raw token IDs only.
+    tokenizer: Option<SentencePieceProcessor>,
     #[allow(dead_code)]
     tokenizer_path: Option<PathBuf>,
 }
@@ -252,14 +250,38 @@ impl Stt {
             lm,
         )
         .map_err(|e| CsmError::Config(format!("moshi::asr::State::new: {e}")))?;
+        let tokenizer_obj = match tokenizer {
+            Some(p) => Some(
+                SentencePieceProcessor::open(p)
+                    .map_err(|e| CsmError::Config(format!("sentencepiece open: {e}")))?,
+            ),
+            None => None,
+        };
         Ok(Self {
             state,
             device: device.clone(),
             pending: Vec::new(),
+            tokenizer: tokenizer_obj,
             tokenizer_path: tokenizer.map(|p| p.to_path_buf()),
         })
     }
 
+    /// Detokenize a Word event's token IDs to text. Skips padding tokens
+    /// (id 3) and uses sentencepiece's built-in detok if a tokenizer was
+    /// loaded.
+    pub fn decode_word_text(&self, tokens: &[u32]) -> Option<String> {
+        let sp = self.tokenizer.as_ref()?;
+        let filtered: Vec<u32> = tokens
+            .iter()
+            .copied()
+            .filter(|&t| t > PADDING_TOKEN_ID)
+            .collect();
+        if filtered.is_empty() {
+            return Some(String::new());
+        }
+        sp.decode_piece_ids(&filtered).ok()
+    }
+
     /// Reset the streaming state for a new utterance/session.
     pub fn reset(&mut self) -> Result<()> {
         self.state