337 lines
8.7 KiB
Rust
337 lines
8.7 KiB
Rust
//! 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(())
|
|
}
|
|
}
|