Initial commit

crates/models/rtx-nlg/src/quality/mod.rs

@@ -0,0 +1,336 @@
+//! Generation quality features and filters
+
+pub mod diversity;
+pub mod repetition;
+pub mod safety;
+
+use crate::Result;
+use serde::{Deserialize, Serialize};
+
+/// Quality filter interface
+pub trait QualityFilter: Send + Sync {
+    /// Check if text passes the quality filter
+    fn passes_filter(&self, text: &str) -> Result<bool>;
+
+    /// Get filter score (0.0 to 1.0)
+    fn filter_score(&self, text: &str) -> Result<f32>;
+
+    /// Get filter name for logging
+    fn name(&self) -> &str;
+}
+
+/// Repetition penalty configuration and implementation
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct RepetitionPenalty {
+    pub penalty_factor: f32,
+    pub ngram_size: usize,
+    pub max_repetitions: usize,
+}
+
+impl Default for RepetitionPenalty {
+    fn default() -> Self {
+        Self {
+            penalty_factor: 1.2,
+            ngram_size: 3,
+            max_repetitions: 2,
+        }
+    }
+}
+
+impl RepetitionPenalty {
+    pub fn new(penalty_factor: f32) -> Self {
+        Self {
+            penalty_factor,
+            ..Default::default()
+        }
+    }
+
+    pub fn calculate_penalty(&self, tokens: &[u32]) -> f32 {
+        if tokens.len() < self.ngram_size {
+            return 0.0;
+        }
+
+        let mut ngram_counts = std::collections::HashMap::new();
+        let mut total_penalty = 0.0;
+
+        // Count n-grams
+        for window in tokens.windows(self.ngram_size) {
+            *ngram_counts.entry(window.to_vec()).or_insert(0) += 1;
+        }
+
+        // Apply penalties for repeated n-grams
+        for (_, count) in ngram_counts {
+            if count > self.max_repetitions {
+                total_penalty += (count - self.max_repetitions) as f32 * self.penalty_factor;
+            }
+        }
+
+        total_penalty
+    }
+}
+
+impl QualityFilter for RepetitionPenalty {
+    fn passes_filter(&self, text: &str) -> Result<bool> {
+        let score = self.filter_score(text)?;
+        Ok(score < 0.5) // Pass if low repetition
+    }
+
+    fn filter_score(&self, text: &str) -> Result<f32> {
+        // Mock tokenization for demonstration
+        let words: Vec<&str> = text.split_whitespace().collect();
+        let tokens: Vec<u32> = words.iter().enumerate().map(|(i, _)| i as u32).collect();
+
+        let penalty = self.calculate_penalty(&tokens);
+        Ok((penalty / words.len() as f32).clamp(0.0, 1.0))
+    }
+
+    fn name(&self) -> &'static str {
+        "repetition_penalty"
+    }
+}
+
+/// Diversity promoter for encouraging varied generation
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct DiversityPromoter {
+    pub min_unique_ratio: f32,
+    pub min_entropy: f32,
+}
+
+impl Default for DiversityPromoter {
+    fn default() -> Self {
+        Self {
+            min_unique_ratio: 0.7,
+            min_entropy: 2.0,
+        }
+    }
+}
+
+impl QualityFilter for DiversityPromoter {
+    fn passes_filter(&self, text: &str) -> Result<bool> {
+        let unique_ratio = self.calculate_unique_ratio(text);
+        let entropy = self.calculate_entropy(text);
+
+        Ok(unique_ratio >= self.min_unique_ratio && entropy >= self.min_entropy)
+    }
+
+    fn filter_score(&self, text: &str) -> Result<f32> {
+        let unique_ratio = self.calculate_unique_ratio(text);
+        let entropy = self.calculate_entropy(text);
+
+        let diversity_score = (unique_ratio * 0.5) + ((entropy / 5.0).clamp(0.0, 1.0) * 0.5);
+        Ok(diversity_score)
+    }
+
+    fn name(&self) -> &'static str {
+        "diversity_promoter"
+    }
+}
+
+impl DiversityPromoter {
+    fn calculate_unique_ratio(&self, text: &str) -> f32 {
+        let words: Vec<&str> = text.split_whitespace().collect();
+        if words.is_empty() {
+            return 0.0;
+        }
+
+        let unique_words: std::collections::HashSet<_> = words.iter().collect();
+        unique_words.len() as f32 / words.len() as f32
+    }
+
+    fn calculate_entropy(&self, text: &str) -> f32 {
+        let words: Vec<&str> = text.split_whitespace().collect();
+        if words.is_empty() {
+            return 0.0;
+        }
+
+        let mut word_counts = std::collections::HashMap::new();
+        for word in &words {
+            *word_counts.entry(*word).or_insert(0) += 1;
+        }
+
+        let total_words = words.len() as f32;
+        let entropy: f32 = word_counts
+            .values()
+            .map(|&count| {
+                let p = count as f32 / total_words;
+                -p * p.log2()
+            })
+            .sum();
+
+        entropy
+    }
+}
+
+/// Toxicity filter using simple keyword detection
+#[derive(Debug, Clone)]
+pub struct ToxicityFilter {
+    toxic_keywords: Vec<String>,
+    threshold: f32,
+}
+
+impl Default for ToxicityFilter {
+    fn default() -> Self {
+        Self {
+            toxic_keywords: vec![
+                "hate".to_string(),
+                "violence".to_string(),
+                // Would include more comprehensive list
+            ],
+            threshold: 0.5,
+        }
+    }
+}
+
+impl ToxicityFilter {
+    pub fn new(threshold: f32) -> Self {
+        Self {
+            threshold,
+            ..Default::default()
+        }
+    }
+}
+
+impl QualityFilter for ToxicityFilter {
+    fn passes_filter(&self, text: &str) -> Result<bool> {
+        let score = self.filter_score(text)?;
+        Ok(score < self.threshold)
+    }
+
+    fn filter_score(&self, text: &str) -> Result<f32> {
+        let text_lower = text.to_lowercase();
+        let mut toxicity_score = 0.0f32;
+
+        for keyword in &self.toxic_keywords {
+            if text_lower.contains(keyword) {
+                toxicity_score += 0.3; // Simple scoring
+            }
+        }
+
+        Ok(toxicity_score.clamp(0.0, 1.0))
+    }
+
+    fn name(&self) -> &'static str {
+        "toxicity_filter"
+    }
+}
+
+/// Combined quality checker
+pub struct QualityChecker {
+    filters: Vec<Box<dyn QualityFilter>>,
+}
+
+impl Default for QualityChecker {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl QualityChecker {
+    pub fn new() -> Self {
+        Self {
+            filters: Vec::new(),
+        }
+    }
+
+    pub fn add_filter(mut self, filter: Box<dyn QualityFilter>) -> Self {
+        self.filters.push(filter);
+        self
+    }
+
+    pub fn check_quality(&self, text: &str) -> Result<QualityReport> {
+        let mut passed_filters = Vec::new();
+        let mut failed_filters = Vec::new();
+        let mut scores = std::collections::HashMap::new();
+
+        for filter in &self.filters {
+            let score = filter.filter_score(text)?;
+            let passes = filter.passes_filter(text)?;
+
+            scores.insert(filter.name().to_string(), score);
+
+            if passes {
+                passed_filters.push(filter.name().to_string());
+            } else {
+                failed_filters.push(filter.name().to_string());
+            }
+        }
+
+        let overall_score = if scores.is_empty() {
+            1.0
+        } else {
+            scores.values().sum::<f32>() / scores.len() as f32
+        };
+
+        Ok(QualityReport {
+            overall_score,
+            passed_filters,
+            failed_filters: failed_filters.clone(),
+            scores,
+            passes_all: failed_filters.is_empty(),
+        })
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct QualityReport {
+    pub overall_score: f32,
+    pub passed_filters: Vec<String>,
+    pub failed_filters: Vec<String>,
+    pub scores: std::collections::HashMap<String, f32>,
+    pub passes_all: bool,
+}
+
+// Re-exports
+pub use diversity::*;
+pub use repetition::*;
+pub use safety::*;
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_repetition_penalty() -> Result<()> {
+        let penalty = RepetitionPenalty::default();
+
+        let repetitive_text = "hello hello hello world world world";
+        let normal_text = "this is a normal sentence without repetition";
+
+        let rep_score = penalty.filter_score(repetitive_text)?;
+        let normal_score = penalty.filter_score(normal_text)?;
+
+        assert!(rep_score > normal_score);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_diversity_promoter() -> Result<()> {
+        let promoter = DiversityPromoter::default();
+
+        let diverse_text = "the quick brown fox jumps over lazy dog";
+        let monotonous_text = "same same same same same same same same";
+
+        let diverse_score = promoter.filter_score(diverse_text)?;
+        let mono_score = promoter.filter_score(monotonous_text)?;
+
+        assert!(diverse_score > mono_score);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_quality_checker() -> Result<()> {
+        let checker = QualityChecker::new()
+            .add_filter(Box::new(RepetitionPenalty::default()))
+            .add_filter(Box::new(DiversityPromoter::default()));
+
+        let text = "This is a diverse and interesting text without repetition";
+        let report = checker.check_quality(text)?;
+
+        assert!(report.overall_score >= 0.0 && report.overall_score <= 1.0);
+        assert_eq!(report.scores.len(), 2);
+
+        Ok(())
+    }
+}