kreuzberg 4.0.0.pre.rc.6

data/vendor/kreuzberg/src/keywords/rake.rs

@@ -0,0 +1,293 @@
+//! RAKE (Rapid Automatic Keyword Extraction) backend implementation.
+
+use super::config::{KeywordConfig, RakeParams};
+use super::types::{Keyword, KeywordAlgorithm};
+use crate::Result;
+use crate::stopwords::{STOPWORDS, get_stopwords};
+use rake::*;
+use std::collections::HashSet;
+
+/// Extract keywords using RAKE algorithm.
+///
+/// RAKE is a co-occurrence based keyword extraction method that:
+/// - Identifies candidate keywords using stop words as delimiters
+/// - Calculates word scores based on frequency and degree
+/// - Combines scores for multi-word phrases
+///
+/// # Arguments
+///
+/// * `text` - The text to extract keywords from
+/// * `config` - Keyword extraction configuration
+///
+/// # Returns
+///
+/// A vector of keywords sorted by relevance (highest score first).
+///
+/// # Errors
+///
+/// Returns an error if keyword extraction fails.
+pub fn extract_keywords_rake(text: &str, config: &KeywordConfig) -> Result<Vec<Keyword>> {
+    let params = config.rake_params.as_ref().cloned().unwrap_or_default();
+    let lang = config.language.as_deref().unwrap_or("en");
+    let normalized_lang = normalize_language_code(lang);
+
+    let mut keywords = run_rake_with_stopwords(text, config, &params, stopwords_for_language(lang));
+
+    if keywords.is_empty() && normalized_lang != "en" {
+        keywords = run_rake_with_stopwords(text, config, &params, stopwords_for_language("en"));
+    }
+
+    Ok(keywords)
+}
+
+fn stopwords_for_language(lang: &str) -> StopWords {
+    if let Some(words) = get_stopwords(lang) {
+        let set: HashSet<String> = words.iter().cloned().collect();
+        return StopWords::from(set);
+    }
+
+    let english: HashSet<String> = STOPWORDS
+        .get("en")
+        .map(|set| set.iter().cloned().collect())
+        .unwrap_or_default();
+    StopWords::from(english)
+}
+
+fn run_rake_with_stopwords(
+    text: &str,
+    config: &KeywordConfig,
+    params: &RakeParams,
+    stopwords: StopWords,
+) -> Vec<Keyword> {
+    let rake = Rake::new(stopwords);
+    let results = rake.run(text);
+
+    let filtered_results: Vec<_> = results
+        .into_iter()
+        .filter_map(|keyword_score| {
+            let keyword = keyword_score.keyword.clone();
+
+            if keyword.len() < params.min_word_length {
+                return None;
+            }
+
+            let word_count = keyword.split_whitespace().count();
+            if word_count > params.max_words_per_phrase {
+                return None;
+            }
+
+            if word_count < config.ngram_range.0 || word_count > config.ngram_range.1 {
+                return None;
+            }
+
+            Some((keyword, keyword_score.score))
+        })
+        .collect();
+
+    if filtered_results.is_empty() {
+        return Vec::new();
+    }
+
+    let min_score = filtered_results.iter().map(|(_, s)| *s).fold(f64::INFINITY, f64::min);
+    let max_score = filtered_results
+        .iter()
+        .map(|(_, s)| *s)
+        .fold(f64::NEG_INFINITY, f64::max);
+
+    let mut keywords: Vec<_> = filtered_results
+        .into_iter()
+        .map(|(keyword, raw_score)| {
+            let normalized_score = if max_score > min_score {
+                ((raw_score - min_score) / (max_score - min_score)).clamp(0.0, 1.0)
+            } else {
+                1.0
+            };
+
+            Keyword::new(keyword, normalized_score as f32, KeywordAlgorithm::Rake)
+        })
+        .collect();
+
+    if config.min_score > 0.0 {
+        keywords.retain(|k| k.score >= config.min_score);
+    }
+
+    keywords.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal));
+
+    keywords.truncate(config.max_keywords);
+
+    keywords
+}
+
+fn normalize_language_code(lang: &str) -> String {
+    let lower = lang.to_lowercase();
+    if let Some(pos) = lower.find(['-', '_']) {
+        lower[..pos].to_string()
+    } else if lower.len() >= 2 {
+        lower[..2].to_string()
+    } else {
+        lower
+    }
+}
+
+#[cfg(all(test, feature = "keywords"))]
+mod tests {
+    use super::*;
+    use crate::keywords::config::RakeParams;
+
+    #[test]
+    fn test_rake_extraction_basic() {
+        let text = "Rust is a systems programming language. \
+                    Rust provides memory safety and performance. \
+                    Memory safety is achieved without garbage collection.";
+
+        let config = KeywordConfig::rake();
+
+        let keywords = extract_keywords_rake(text, &config).unwrap();
+
+        assert!(!keywords.is_empty(), "Should extract keywords");
+        assert!(
+            keywords.len() <= config.max_keywords,
+            "Should respect max_keywords limit"
+        );
+
+        for i in 1..keywords.len() {
+            assert!(
+                keywords[i - 1].score >= keywords[i].score,
+                "Keywords should be sorted by score"
+            );
+        }
+
+        for keyword in &keywords {
+            assert_eq!(keyword.algorithm, KeywordAlgorithm::Rake);
+        }
+    }
+
+    #[test]
+    fn test_rake_extraction_with_min_score() {
+        let text = "Rust programming language provides memory safety without garbage collection.";
+
+        let config = KeywordConfig::rake().with_min_score(0.3);
+
+        let keywords = extract_keywords_rake(text, &config).unwrap();
+
+        for keyword in &keywords {
+            assert!(
+                keyword.score >= config.min_score,
+                "Keyword score {} should be >= min_score {}",
+                keyword.score,
+                config.min_score
+            );
+        }
+    }
+
+    #[test]
+    fn test_rake_extraction_with_ngram_range() {
+        let text = "Machine learning models require large datasets for training.";
+
+        let config = KeywordConfig::rake().with_ngram_range(1, 1);
+        let keywords = extract_keywords_rake(text, &config).unwrap();
+
+        for keyword in &keywords {
+            assert_eq!(
+                keyword.text.split_whitespace().count(),
+                1,
+                "Should only extract unigrams"
+            );
+        }
+    }
+
+    #[test]
+    fn test_rake_extraction_empty_text() {
+        let config = KeywordConfig::rake();
+        let keywords = extract_keywords_rake("", &config).unwrap();
+        assert!(keywords.is_empty(), "Empty text should yield no keywords");
+    }
+
+    #[test]
+    fn test_rake_extraction_with_custom_params() {
+        let text = "Natural language processing enables computers to understand human language.";
+
+        let params = RakeParams {
+            min_word_length: 3,
+            max_words_per_phrase: 2,
+        };
+
+        let config = KeywordConfig::rake().with_rake_params(params);
+
+        let keywords = extract_keywords_rake(text, &config).unwrap();
+
+        for keyword in &keywords {
+            for word in keyword.text.split_whitespace() {
+                assert!(word.len() >= 3, "Word '{}' should have min length 3", word);
+            }
+        }
+
+        for keyword in &keywords {
+            assert!(
+                keyword.text.split_whitespace().count() <= 2,
+                "Keyword '{}' should have max 2 words",
+                keyword.text
+            );
+        }
+    }
+
+    #[cfg_attr(coverage, ignore = "coverage instrumentation interferes with stopword heuristics")]
+    #[test]
+    fn test_rake_multilingual() {
+        let spanish_text = "El idioma español es una lengua romance.";
+        let config = KeywordConfig::rake().with_language("es");
+        let keywords = extract_keywords_rake(spanish_text, &config).unwrap();
+        assert!(!keywords.is_empty(), "Should extract Spanish keywords");
+
+        assert!(
+            keywords
+                .iter()
+                .any(|k| k.text.contains("idioma") || k.text.contains("español") || k.text.contains("lengua"))
+        );
+
+        let english_text = "Natural language processing is a subfield of artificial intelligence.";
+        let config = KeywordConfig::rake().with_language("fr");
+        let keywords = extract_keywords_rake(english_text, &config).unwrap();
+        assert!(
+            !keywords.is_empty(),
+            "Should fall back to English stopwords and extract keywords"
+        );
+    }
+
+    #[test]
+    fn test_rake_score_normalization() {
+        let text = "Rust is a systems programming language that provides memory safety and \
+                    thread safety without garbage collection. Rust uses a ownership system \
+                    with rules that the compiler checks at compile time.";
+
+        let config = KeywordConfig::rake();
+        let keywords = extract_keywords_rake(text, &config).unwrap();
+
+        assert!(!keywords.is_empty(), "Should extract keywords");
+
+        for keyword in &keywords {
+            assert!(
+                keyword.score >= 0.0 && keyword.score <= 1.0,
+                "Keyword '{}' score {} should be in range [0.0, 1.0]",
+                keyword.text,
+                keyword.score
+            );
+        }
+
+        if !keywords.is_empty() {
+            let max_score = keywords.iter().map(|k| k.score).fold(0.0f32, f32::max);
+            assert!(
+                (max_score - 1.0).abs() < 0.001,
+                "Max score should be 1.0, got {}",
+                max_score
+            );
+        }
+
+        for i in 1..keywords.len() {
+            assert!(
+                keywords[i - 1].score >= keywords[i].score,
+                "Keywords should be sorted by score"
+            );
+        }
+    }
+}

data/vendor/kreuzberg/src/keywords/types.rs

@@ -0,0 +1,68 @@
+//! Shared types for keyword extraction.
+
+use serde::{Deserialize, Serialize};
+
+/// Keyword algorithm selection.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum KeywordAlgorithm {
+    /// YAKE (Yet Another Keyword Extractor) - statistical approach
+    #[cfg(feature = "keywords-yake")]
+    Yake,
+
+    /// RAKE (Rapid Automatic Keyword Extraction) - co-occurrence based
+    #[cfg(feature = "keywords-rake")]
+    Rake,
+}
+
+impl Default for KeywordAlgorithm {
+    fn default() -> Self {
+        #[cfg(feature = "keywords-yake")]
+        return Self::Yake;
+
+        #[cfg(all(feature = "keywords-rake", not(feature = "keywords-yake")))]
+        return Self::Rake;
+
+        #[cfg(not(any(feature = "keywords-yake", feature = "keywords-rake")))]
+        compile_error!("At least one keyword extraction feature must be enabled");
+    }
+}
+
+/// Extracted keyword with metadata.
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct Keyword {
+    /// The keyword text.
+    pub text: String,
+
+    /// Relevance score (higher is better, algorithm-specific range).
+    pub score: f32,
+
+    /// Algorithm that extracted this keyword.
+    pub algorithm: KeywordAlgorithm,
+
+    /// Optional positions where keyword appears in text (character offsets).
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub positions: Option<Vec<usize>>,
+}
+
+impl Keyword {
+    /// Create a new keyword.
+    pub fn new(text: String, score: f32, algorithm: KeywordAlgorithm) -> Self {
+        Self {
+            text,
+            score,
+            algorithm,
+            positions: None,
+        }
+    }
+
+    /// Create a new keyword with positions.
+    pub fn with_positions(text: String, score: f32, algorithm: KeywordAlgorithm, positions: Vec<usize>) -> Self {
+        Self {
+            text,
+            score,
+            algorithm,
+            positions: Some(positions),
+        }
+    }
+}

data/vendor/kreuzberg/src/keywords/yake.rs

@@ -0,0 +1,163 @@
+//! YAKE (Yet Another Keyword Extractor) backend implementation.
+
+use super::config::KeywordConfig;
+use super::types::{Keyword, KeywordAlgorithm};
+use crate::Result;
+use yake_rust::{Config as YakeConfig, StopWords, get_n_best};
+
+/// Extract keywords using YAKE algorithm.
+///
+/// YAKE is a statistical keyword extraction method that weighs multiple factors:
+/// - Term frequency and position
+/// - Capitalization
+/// - Sentence co-occurrence
+/// - Context analysis
+///
+/// # Arguments
+///
+/// * `text` - The text to extract keywords from
+/// * `config` - Keyword extraction configuration
+///
+/// # Returns
+///
+/// A vector of keywords sorted by relevance (highest score first).
+///
+/// # Errors
+///
+/// Returns an error if keyword extraction fails.
+pub fn extract_keywords_yake(text: &str, config: &KeywordConfig) -> Result<Vec<Keyword>> {
+    let params = config.yake_params.as_ref().cloned().unwrap_or_default();
+
+    let yake_config = YakeConfig {
+        ngrams: config.ngram_range.1,
+        window_size: params.window_size,
+        ..YakeConfig::default()
+    };
+
+    let stopwords = if let Some(ref lang) = config.language {
+        StopWords::predefined(lang).unwrap_or_else(|| {
+            tracing::debug!(
+                "WARNING: Stopwords not available for language '{}', using default English stopwords",
+                lang
+            );
+            StopWords::default()
+        })
+    } else {
+        StopWords::default()
+    };
+
+    let results = get_n_best(config.max_keywords, text, &stopwords, &yake_config);
+
+    let mut keywords = results
+        .into_iter()
+        .filter(|item| {
+            let word_count = item.keyword.split_whitespace().count();
+            word_count >= config.ngram_range.0
+        })
+        .map(|item| {
+            let normalized_score = if item.score > 0.0 {
+                (1.0_f64 / (1.0 + item.score)).clamp(0.0, 1.0)
+            } else {
+                1.0
+            };
+
+            Keyword::new(item.keyword, normalized_score as f32, KeywordAlgorithm::Yake)
+        })
+        .collect::<Vec<_>>();
+
+    if config.min_score > 0.0 {
+        keywords.retain(|k| k.score >= config.min_score);
+    }
+
+    keywords.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal));
+
+    Ok(keywords)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::keywords::config::YakeParams;
+
+    #[test]
+    fn test_yake_extraction_basic() {
+        let text = "Rust is a systems programming language. \
+                    Rust provides memory safety and performance. \
+                    Memory safety is achieved without garbage collection.";
+
+        let config = KeywordConfig::yake();
+
+        let keywords = extract_keywords_yake(text, &config).unwrap();
+
+        assert!(!keywords.is_empty(), "Should extract keywords");
+        assert!(
+            keywords.len() <= config.max_keywords,
+            "Should respect max_keywords limit"
+        );
+
+        for i in 1..keywords.len() {
+            assert!(
+                keywords[i - 1].score >= keywords[i].score,
+                "Keywords should be sorted by score"
+            );
+        }
+
+        for keyword in &keywords {
+            assert_eq!(keyword.algorithm, KeywordAlgorithm::Yake);
+        }
+    }
+
+    #[test]
+    fn test_yake_extraction_with_min_score() {
+        let text = "Rust programming language provides memory safety without garbage collection.";
+
+        let config = KeywordConfig::yake().with_min_score(0.5);
+
+        let keywords = extract_keywords_yake(text, &config).unwrap();
+
+        for keyword in &keywords {
+            assert!(
+                keyword.score >= config.min_score,
+                "Keyword score {} should be >= min_score {}",
+                keyword.score,
+                config.min_score
+            );
+        }
+    }
+
+    #[test]
+    fn test_yake_extraction_with_ngram_range() {
+        let text = "Machine learning models require large datasets for training.";
+
+        let config = KeywordConfig::yake().with_ngram_range(1, 1);
+        let keywords = extract_keywords_yake(text, &config).unwrap();
+
+        for keyword in &keywords {
+            assert_eq!(
+                keyword.text.split_whitespace().count(),
+                1,
+                "Should only extract unigrams"
+            );
+        }
+    }
+
+    #[test]
+    fn test_yake_extraction_empty_text() {
+        let config = KeywordConfig::yake();
+        let keywords = extract_keywords_yake("", &config).unwrap();
+        assert!(keywords.is_empty(), "Empty text should yield no keywords");
+    }
+
+    #[test]
+    fn test_yake_extraction_with_custom_params() {
+        let text = "Natural language processing enables computers to understand human language.";
+
+        let params = YakeParams { window_size: 3 };
+
+        let config = KeywordConfig::yake().with_yake_params(params);
+
+        let keywords = extract_keywords_yake(text, &config).unwrap();
+
+        assert!(!keywords.is_empty(), "Should extract keywords with custom params");
+    }
+}