rustytorch/crates/models/rtx-csm/src/repetition.rs

//! Frame-level repetition detection ("loop-escape").
//!
//! CSM's known pathological mode is **syllable looping** — the AR decoder
//! gets stuck in a low-entropy attractor and emits the same 32-codebook
//! frame over and over. Users hear it as a syllable repeating forever.
//!
//! The ideal fix is per-codebook logit blocking on recent frame
//! fingerprints, but the actual sampling happens *inside* candle's
//! `csm::Model::generate_frame` and we can't hook between codebook
//! samples without forking the upstream module. So instead we detect the
//! loop *after* the frame is emitted and short-circuit generation.
//!
//! This catches both:
//!   - period-1 loops (frame_k == frame_{k-1}, ...)
//!   - higher-period loops where a small set of frames cycles
//!
//! Detection rule: within the last `window` emitted frames, if any single
//! fingerprint occurs `max_repeats` or more times, we declare a loop.

use std::collections::{HashMap, VecDeque};

#[derive(Debug, Clone, Copy)]
pub struct RepetitionConfig {
    /// How many recent frames to track. 16 frames = 1.28s at 12.5 Hz.
    pub window: usize,
    /// A single fingerprint occurring this many times in the window triggers
    /// a loop break. 4 is a safe default — natural speech can have 2-3
    /// repeated frames at a sustained vowel without being pathological.
    pub max_repeats: usize,
}

impl Default for RepetitionConfig {
    fn default() -> Self {
        Self {
            window: 16,
            max_repeats: 4,
        }
    }
}

/// Tracks recent frame fingerprints. Cheap (a small ring buffer of u64).
pub struct RepetitionGuard {
    cfg: RepetitionConfig,
    recent: VecDeque<u64>,
}

impl RepetitionGuard {
    pub fn new(cfg: RepetitionConfig) -> Self {
        Self {
            cfg,
            recent: VecDeque::with_capacity(cfg.window),
        }
    }

    pub fn reset(&mut self) {
        self.recent.clear();
    }

    /// Record a frame and return true if a loop is detected. Frame is the
    /// `Vec<u32>` of length num_codebooks returned by `generate_frame`.
    pub fn observe(&mut self, frame: &[u32]) -> bool {
        let fp = fingerprint(frame);
        if self.recent.len() == self.cfg.window {
            self.recent.pop_front();
        }
        self.recent.push_back(fp);
        // Count occurrences only when the window has enough frames to make
        // the threshold meaningful — avoids spurious early triggers.
        if self.recent.len() < self.cfg.max_repeats {
            return false;
        }
        let mut counts: HashMap<u64, usize> = HashMap::new();
        let mut max_count = 0usize;
        for h in self.recent.iter() {
            let c = counts.entry(*h).or_insert(0);
            *c += 1;
            if *c > max_count {
                max_count = *c;
            }
        }
        max_count >= self.cfg.max_repeats
    }
}

/// FNV-1a 64-bit hash over the codebook ids. Fast, avoids extra deps.
fn fingerprint(frame: &[u32]) -> u64 {
    const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
    const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
    let mut h = FNV_OFFSET;
    for v in frame {
        let bytes = v.to_le_bytes();
        for b in bytes {
            h ^= b as u64;
            h = h.wrapping_mul(FNV_PRIME);
        }
    }
    h
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn no_repeats_no_trigger() {
        let mut g = RepetitionGuard::new(RepetitionConfig::default());
        for i in 0..20u32 {
            let frame: Vec<u32> = (0..32).map(|c| i.wrapping_add(c)).collect();
            assert!(!g.observe(&frame), "false trigger at frame {i}");
        }
    }

    #[test]
    fn period_1_loop_trips_after_max_repeats() {
        let mut g = RepetitionGuard::new(RepetitionConfig {
            window: 16,
            max_repeats: 4,
        });
        let frame = vec![42u32; 32];
        // Three identical frames must NOT trigger.
        assert!(!g.observe(&frame));
        assert!(!g.observe(&frame));
        assert!(!g.observe(&frame));
        // Fourth one trips the guard.
        assert!(g.observe(&frame));
    }

    #[test]
    fn period_2_alternating_loop_trips() {
        let mut g = RepetitionGuard::new(RepetitionConfig {
            window: 16,
            max_repeats: 4,
        });
        let a = vec![1u32; 32];
        let b = vec![2u32; 32];
        let mut tripped_at = None;
        for i in 0..16 {
            let f = if i % 2 == 0 { &a } else { &b };
            if g.observe(f) {
                tripped_at = Some(i);
                break;
            }
        }
        assert!(tripped_at.is_some(), "alternating loop never tripped");
    }

    #[test]
    fn fingerprint_distinguishes_neighbors() {
        let a: Vec<u32> = (0..32).collect();
        let mut b = a.clone();
        b[5] = 999;
        assert_ne!(fingerprint(&a), fingerprint(&b));
    }

    #[test]
    fn reset_clears_history() {
        let mut g = RepetitionGuard::new(RepetitionConfig {
            window: 16,
            max_repeats: 4,
        });
        let frame = vec![42u32; 32];
        for _ in 0..5 {
            g.observe(&frame);
        }
        g.reset();
        // After reset, we need 4 more observations to trip again.
        assert!(!g.observe(&frame));
        assert!(!g.observe(&frame));
        assert!(!g.observe(&frame));
        assert!(g.observe(&frame));
    }
}