context_forge/analysis/
ngrams.rs

1/// Generate n-grams from a token list.
2///
3/// Produces n-grams by sliding a window of size `n` over the tokens
4/// and joining with a single space.
5///
6/// Returns an empty vec if `tokens.len()` < n.
7#[must_use]
8pub fn extract(tokens: &[String], n: usize) -> Vec<String> {
9    if n == 0 || tokens.len() < n {
10        return Vec::new();
11    }
12
13    tokens.windows(n).map(|window| window.join(" ")).collect()
14}
15
16/// Generate bigrams from a token list.
17#[must_use]
18pub fn bigrams(tokens: &[String]) -> Vec<String> {
19    extract(tokens, 2)
20}
21
22/// Generate trigrams from a token list.
23#[must_use]
24pub fn trigrams(tokens: &[String]) -> Vec<String> {
25    extract(tokens, 3)
26}
27
28#[cfg(test)]
29mod tests {
30    use super::{bigrams, extract, trigrams};
31
32    fn sample_tokens() -> Vec<String> {
33        ["a", "b", "c", "d"]
34            .into_iter()
35            .map(str::to_string)
36            .collect()
37    }
38
39    #[test]
40    fn test_bigrams() {
41        let tokens = sample_tokens();
42        let expected = vec!["a b".to_string(), "b c".to_string(), "c d".to_string()];
43        assert_eq!(bigrams(&tokens), expected);
44    }
45
46    #[test]
47    fn test_trigrams() {
48        let tokens = sample_tokens();
49        let expected = vec!["a b c".to_string(), "b c d".to_string()];
50        assert_eq!(trigrams(&tokens), expected);
51    }
52
53    #[test]
54    fn test_ngrams_insufficient_tokens() {
55        let tokens = vec!["a".to_string()];
56        assert!(extract(&tokens, 2).is_empty());
57    }
58
59    #[test]
60    fn test_ngrams_exact_size() {
61        let tokens = vec!["a".to_string(), "b".to_string()];
62        assert_eq!(extract(&tokens, 2), vec!["a b".to_string()]);
63    }
64
65    #[test]
66    fn test_unigrams() {
67        let tokens = sample_tokens();
68        assert_eq!(extract(&tokens, 1), tokens);
69    }
70
71    #[test]
72    fn test_zero_gram() {
73        let tokens = sample_tokens();
74        assert!(extract(&tokens, 0).is_empty());
75    }
76}
context_forge/analysis/ngrams.rs

context_forge/analysis/
ngrams.rs
