Files
rustytorch/crates/models/rtx-nlg/src/quality/repetition.rs
T
2026-03-04 00:08:42 +00:00

113 lines
3.0 KiB
Rust

//! Repetition detection and penalty implementation
use std::collections::HashMap;
/// Advanced repetition detection
pub struct RepetitionDetector {
ngram_size: usize,
window_size: usize,
}
impl RepetitionDetector {
pub fn new(ngram_size: usize, window_size: usize) -> Self {
Self {
ngram_size,
window_size,
}
}
pub fn detect_repetitions(&self, tokens: &[u32]) -> RepetitionReport {
let mut ngram_counts = HashMap::new();
let mut repetitive_spans = Vec::new();
// Count n-grams in sliding windows
for window_start in 0..tokens.len().saturating_sub(self.window_size) {
let window_end = (window_start + self.window_size).min(tokens.len());
let window = &tokens[window_start..window_end];
for ngram in window.windows(self.ngram_size) {
*ngram_counts.entry(ngram.to_vec()).or_insert(0) += 1;
}
}
// Find repetitive patterns
for (ngram, count) in ngram_counts {
if count > 2 {
repetitive_spans.push(RepetitiveSpan {
ngram,
count,
severity: self.calculate_severity(count),
});
}
}
let overall_score = self.calculate_overall_score(&repetitive_spans);
RepetitionReport {
total_repetitions: repetitive_spans.len(),
repetitive_spans,
overall_score,
}
}
fn calculate_severity(&self, count: usize) -> RepetitionSeverity {
match count {
3..=5 => RepetitionSeverity::Mild,
6..=10 => RepetitionSeverity::Moderate,
_ => RepetitionSeverity::Severe,
}
}
fn calculate_overall_score(&self, spans: &[RepetitiveSpan]) -> f32 {
if spans.is_empty() {
return 1.0;
}
let severity_sum: f32 = spans
.iter()
.map(|span| match span.severity {
RepetitionSeverity::Mild => 0.1,
RepetitionSeverity::Moderate => 0.3,
RepetitionSeverity::Severe => 0.5,
})
.sum();
(1.0 - severity_sum / spans.len() as f32).clamp(0.0, 1.0)
}
}
#[derive(Debug, Clone)]
pub struct RepetitionReport {
pub total_repetitions: usize,
pub repetitive_spans: Vec<RepetitiveSpan>,
pub overall_score: f32,
}
#[derive(Debug, Clone)]
pub struct RepetitiveSpan {
pub ngram: Vec<u32>,
pub count: usize,
pub severity: RepetitionSeverity,
}
#[derive(Debug, Clone, PartialEq)]
pub enum RepetitionSeverity {
Mild,
Moderate,
Severe,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_repetition_detection() {
let detector = RepetitionDetector::new(3, 20);
let tokens = vec![1, 2, 3, 1, 2, 3, 1, 2, 3, 4, 5, 6];
let report = detector.detect_repetitions(&tokens);
assert!(report.total_repetitions > 0);
assert!(report.overall_score < 1.0);
}
}