N-gram Language Model

Status: Verified | Idempotent: Yes | Coverage: 95%+

Build a simple n-gram language model for text generation.

Run Command

cargo run --example create_apr_ngram_language_model

Code

//! # Recipe: Create APR N-gram Language Model
//!
//! Contract: contracts/recipe-iiur-v1.yaml
//! **Category**: Model Creation
//! **Isolation Level**: Full
//! **Idempotency**: Guaranteed
//! **Dependencies**: None (default features)
//!
//! ## QA Checklist
//! 1. [x] `cargo run` succeeds (Exit Code 0)
//! 2. [x] `cargo test` passes
//! 3. [x] Deterministic output (Verified)
//! 4. [x] No temp files leaked
//! 5. [x] Memory usage stable
//! 6. [x] WASM compatible (N/A)
//! 7. [x] Clippy clean
//! 8. [x] Rustfmt standard
//! 9. [x] No `unwrap()` in logic
//! 10. [x] Proptests pass (100+ cases)
//!
//! ## Learning Objective
//! Build an N-gram language model from a text corpus and save as `.apr`.
//!
//! ## Run Command
//! ```bash
//! cargo run --example create_apr_ngram_language_model
//! ```
//!
//!
//! ## Format Variants
//! ```bash
//! apr convert model.apr          # APR native format
//! apr convert model.gguf         # GGUF (llama.cpp compatible)
//! apr convert model.safetensors  # SafeTensors (HuggingFace)
//! ```
//! ## References
//! - Jacob, B. et al. (2018). *Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference*. CVPR. arXiv:1712.05877

use apr_cookbook::prelude::*;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

fn main() -> Result<()> {
    let mut ctx = RecipeContext::new("create_apr_ngram_language_model")?;

    // Sample corpus for training
    let corpus = [
        "the quick brown fox jumps over the lazy dog",
        "the quick brown fox runs through the forest",
        "a lazy dog sleeps in the sun",
        "the brown dog chases the quick fox",
        "quick thinking leads to quick results",
    ];

    // Build N-gram model
    let n = 3; // Trigram model
    let model = build_ngram_model(&corpus, n);

    ctx.record_metric("n", n as i64);
    ctx.record_metric("vocabulary_size", model.vocabulary.len() as i64);
    ctx.record_metric("ngram_count", model.ngrams.len() as i64);

    // Test generation
    let seed_words = vec!["the".to_string(), "quick".to_string()];
    let generated = generate_text(&model, &seed_words, 10);
    ctx.record_string_metric("generated_sample", generated.join(" "));

    // Serialize and save
    let model_bytes = serialize_ngram_model(&model)?;

    let mut converter = AprConverter::new();
    converter.set_metadata(ConversionMetadata {
        name: Some("ngram-lm".to_string()),
        architecture: Some("ngram".to_string()),
        source_format: None,
        custom: HashMap::new(),
    });

    converter.add_tensor(TensorData {
        name: "ngram_model".to_string(),
        shape: vec![model_bytes.len()],
        dtype: DataType::U8,
        data: model_bytes,
    });

    let apr_path = ctx.path("ngram_lm.apr");
    let apr_bytes = converter.to_apr()?;
    std::fs::write(&apr_path, &apr_bytes)?;

    println!("=== Recipe: {} ===", ctx.name());
    println!("Built {}-gram model", n);
    println!("Vocabulary size: {}", model.vocabulary.len());
    println!("N-gram count: {}", model.ngrams.len());
    println!("Generated text: {}", generated.join(" "));
    println!("Saved to: {:?}", apr_path);

    Ok(())
}

/// N-gram language model
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NgramModel {
    /// N-gram order (2 = bigram, 3 = trigram)
    pub n: usize,
    /// Vocabulary (word -> index)
    pub vocabulary: HashMap<String, usize>,
    /// N-gram counts: (context -> (next_word -> count))
    pub ngrams: HashMap<String, HashMap<String, usize>>,
}

/// Build an N-gram model from a corpus
fn build_ngram_model(corpus: &[&str], n: usize) -> NgramModel {
    let mut vocabulary = HashMap::new();
    let mut ngrams: HashMap<String, HashMap<String, usize>> = HashMap::new();

    for sentence in corpus {
        let words: Vec<&str> = sentence.split_whitespace().collect();

        // Build vocabulary
        for word in &words {
            let idx = vocabulary.len();
            vocabulary.entry((*word).to_string()).or_insert(idx);
        }

        // Extract n-grams
        if words.len() >= n {
            for window in words.windows(n) {
                let context = window[..n - 1].join(" ");
                let next_word = window[n - 1].to_string();

                ngrams
                    .entry(context)
                    .or_default()
                    .entry(next_word)
                    .and_modify(|c| *c += 1)
                    .or_insert(1);
            }
        }
    }

    NgramModel {
        n,
        vocabulary,
        ngrams,
    }
}

/// Generate text using the N-gram model
fn generate_text(model: &NgramModel, seed: &[String], max_words: usize) -> Vec<String> {
    let mut result = seed.to_vec();
    let context_len = model.n - 1;

    for _ in 0..max_words {
        if result.len() < context_len {
            break;
        }

        let context = result[result.len() - context_len..].join(" ");

        match model.ngrams.get(&context) {
            Some(next_words) => {
                // Pick the most likely next word (deterministic for reproducibility)
                if let Some((word, _)) = next_words.iter().max_by_key(|(_, &count)| count) {
                    result.push(word.clone());
                } else {
                    break;
                }
            }
            None => break,
        }
    }

    result
}

/// Calculate perplexity on a test sentence
#[allow(dead_code)]
fn calculate_perplexity(model: &NgramModel, sentence: &str) -> f64 {
    let words: Vec<&str> = sentence.split_whitespace().collect();
    let context_len = model.n - 1;

    if words.len() < model.n {
        return f64::INFINITY;
    }

    let mut log_prob_sum = 0.0f64;
    let mut count = 0;

    for window in words.windows(model.n) {
        let context = window[..context_len].join(" ");
        let next_word = window[context_len];

        let prob = match model.ngrams.get(&context) {
            Some(next_words) => {
                let total: usize = next_words.values().sum();
                let word_count = next_words.get(next_word).copied().unwrap_or(1);
                word_count as f64 / total as f64
            }
            None => 1.0 / model.vocabulary.len() as f64, // Smoothing
        };

        log_prob_sum += prob.ln();
        count += 1;
    }

    if count == 0 {
        return f64::INFINITY;
    }

    (-log_prob_sum / f64::from(count)).exp()
}

fn serialize_ngram_model(model: &NgramModel) -> Result<Vec<u8>> {
    serde_json::to_vec(model).map_err(|e| CookbookError::Serialization(e.to_string()))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_build_model() {
        let corpus = ["a b c", "a b d"];
        let model = build_ngram_model(&corpus, 2);

        assert!(model.vocabulary.contains_key("a"));
        assert!(model.vocabulary.contains_key("b"));
        assert!(model.vocabulary.contains_key("c"));
        assert!(model.vocabulary.contains_key("d"));
        assert_eq!(model.vocabulary.len(), 4);
    }

    #[test]
    fn test_ngram_extraction() {
        let corpus = ["a b c d"];
        let model = build_ngram_model(&corpus, 2);

        // Should have bigrams: "a" -> "b", "b" -> "c", "c" -> "d"
        assert!(model.ngrams.contains_key("a"));
        assert!(model.ngrams.contains_key("b"));
        assert!(model.ngrams.contains_key("c"));
    }

    #[test]
    fn test_trigram_extraction() {
        let corpus = ["a b c d e"];
        let model = build_ngram_model(&corpus, 3);

        // Should have trigrams: "a b" -> "c", "b c" -> "d", "c d" -> "e"
        assert!(model.ngrams.contains_key("a b"));
        assert!(model.ngrams.contains_key("b c"));
        assert!(model.ngrams.contains_key("c d"));
    }

    #[test]
    fn test_text_generation() {
        let corpus = ["the cat sat", "the cat ran", "the dog sat"];
        let model = build_ngram_model(&corpus, 2);

        let seed = vec!["the".to_string()];
        let generated = generate_text(&model, &seed, 5);

        // Should start with seed
        assert_eq!(generated[0], "the");
        // Should generate something after
        assert!(generated.len() > 1);
    }

    #[test]
    fn test_perplexity() {
        let corpus = ["a b c", "a b c"];
        let model = build_ngram_model(&corpus, 2);

        let perp = calculate_perplexity(&model, "a b c");
        assert!(perp.is_finite());
        assert!(perp > 0.0);
    }

    #[test]
    fn test_serialization() {
        let corpus = ["test sentence"];
        let model = build_ngram_model(&corpus, 2);
        let bytes = serialize_ngram_model(&model).unwrap();
        assert!(!bytes.is_empty());
    }

    #[test]
    fn test_deterministic() {
        let corpus = ["a b c d"];

        let model1 = build_ngram_model(&corpus, 2);
        let model2 = build_ngram_model(&corpus, 2);

        // Same corpus should produce same vocabulary
        assert_eq!(model1.vocabulary.len(), model2.vocabulary.len());
        assert_eq!(model1.ngrams.len(), model2.ngrams.len());
    }
}

#[cfg(test)]
mod proptests {
    use super::*;
    use proptest::prelude::*;

    proptest! {
        #![proptest_config(ProptestConfig::with_cases(50))]

        #[test]
        fn prop_vocabulary_size(words in proptest::collection::vec("[a-z]+", 1..20)) {
            let sentence = words.join(" ");
            let corpus = [sentence.as_str()];
            let model = build_ngram_model(&corpus, 2);

            // Vocabulary should be at most the number of unique words
            let unique_words: std::collections::HashSet<_> = words.iter().collect();
            prop_assert!(model.vocabulary.len() <= unique_words.len());
        }

        #[test]
        fn prop_ngram_order(n in 2usize..5) {
            let corpus = ["a b c d e f g h"];
            let model = build_ngram_model(&corpus, n);

            prop_assert_eq!(model.n, n);
        }
    }
}

Key Concepts

1. N-gram Storage: Context-to-next-word probability mappings
2. Vocabulary: Token-to-index mapping stored in model metadata
3. Smoothing: Handle unseen n-grams with backoff strategies