This document describes how data flows through libgrammstein during training and inference.
┌─────────────────────────────────────────────────────────────────┐
│ libgrammstein │
│ │
│ ┌─────────────┐ ┌──────────────┐ ┌───────────────────┐ │
│ │ Corpus │───>│ Training │───>│ Model │ │
│ │ Readers │ │ Pipeline │ │ (Serialized) │ │
│ └─────────────┘ └──────────────┘ └───────────────────┘ │
│ │ │ │ │
│ │ │ │ │
│ ▼ ▼ ▼ │
│ ┌─────────────┐ ┌──────────────┐ ┌───────────────────┐ │
│ │ Quality │ │ Dictionary │ │ Queries │ │
│ │ Filters │ │ Backend │ │ (log_prob, etc) │ │
│ └─────────────┘ └──────────────┘ └───────────────────┘ │
│ │
└─────────────────────────────────────────────────────────────────┘
Raw Text ──> CorpusReader ──> Sentences ──> Quality Filter ──> Clean Sentences
-
Input Sources
- Plain text files (
PlaintextReader) - Wikipedia XML dumps (
WikipediaReader) - Project Gutenberg texts (
GutenbergReader) - HTTP streams (Wikipedia URLs)
- Plain text files (
-
Processing
- Sentence segmentation
- Unicode normalization
- Quality filtering (optional)
- Deduplication (optional)
// Example flow
let reader = WikipediaReader::from_dump("enwiki.xml.bz2")?;
// Sentences are produced lazily
for sentence in reader.sentences() {
// Each sentence is a clean string
process(sentence);
}Clean Sentences ──> Tokenizer ──> Tokens ──> N-gram Extractor ──> Dictionary
-
Tokenization
- Split on whitespace
- Lowercase (optional)
- Remove punctuation (optional)
-
N-gram Extraction
- Extract all n-grams up to order N
- For sentence ["the", "quick", "fox"]:
- Unigrams: "the", "quick", "fox"
- Bigrams: "the quick", "quick fox"
- Trigrams: "the quick fox"
-
Counting
- Atomic increment in dictionary
- Track continuation counts for MKN
// Internal flow
for n in 1..=order {
for i in 0..=(tokens.len() - n) {
let ngram = &tokens[i..i+n];
let key = ngram.join(" ");
dictionary.insert_or_increment(&key);
}
}Raw Counts ──> Continuation Counts ──> Discount Parameters ──> Smoothed Model
-
Continuation Counts
- N₁, N₂, N₃+: Count of n-grams occurring 1, 2, 3+ times
- Used for Modified Kneser-Ney discount calculation
-
Discount Parameters
- D₁ = 1 - 2Y × N₂/N₁
- D₂ = 2 - 3Y × N₃/N₂
- D₃₊ = 3 - 4Y × N₄/N₃
- Where Y = N₁/(N₁ + 2N₂)
Sentences ──> Vocabulary ──> Skip-gram ──> Embeddings
│
▼
Negative Sampling
-
Vocabulary Building
- Count word frequencies
- Filter by min_count
- Build word-to-index mapping
-
Skip-gram Training
- For each word, predict context words
- Update word and subword embeddings
- Use negative sampling for efficiency
Query ──> Context Lookup ──> Backoff Chain ──> Smoothed Probability
// Query: P(fox | quick brown)
let query = ["quick", "brown", "fox"];
// 1. Try trigram "quick brown fox"
if let Some(count) = trie.get("quick brown fox") {
// Compute smoothed probability
return mkn_probability(count, context_count);
}
// 2. Backoff to bigram "brown fox"
if let Some(count) = trie.get("brown fox") {
return backoff_weight * mkn_probability(count, context_count);
}
// 3. Backoff to unigram "fox"
return backoff_weight² * unigram_probability("fox");Word ──> Known? ─yes─> Word Embedding + Subword Average
│
no
│
▼
Subword Extraction ──> Hash to Buckets ──> Average Subword Embeddings
// Query: vector("hello")
if let Some(idx) = word_to_idx.get("hello") {
// Known word: combine word and subword embeddings
let word_vec = word_embeddings.row(idx);
let subword_vec = average_subword_vectors("hello");
return (word_vec + subword_vec) / 2.0;
} else {
// OOV: use only subword embeddings
return average_subword_vectors("hello");
} ┌──────────────┐
│ N-gram │──> log P_ngram
Word + Context ─────────>│ Model │
└──────────────┘
│
├─────────> Interpolate ──> Final Score
│
┌──────────────┐
│ Embedding │──> log P_embed
Word + Context ─────────>│ Model │
└──────────────┘
// Interpolation strategies
match strategy {
Linear { alpha } => {
let p_ngram = ngram.log_prob(word, context).exp();
let p_embed = embedding_prob(word, context).exp();
(alpha * p_ngram + (1-alpha) * p_embed).ln()
}
LogLinear { alpha } => {
let lp_ngram = ngram.log_prob(word, context);
let lp_embed = embedding_prob(word, context);
alpha * lp_ngram + (1-alpha) * lp_embed
}
NgramWithEmbeddingFallback => {
if ngram.in_vocabulary(word) {
ngram.log_prob(word, context)
} else {
embedding_prob(word, context)
}
}
}Insert ──> Find Prefix ──> Add Nodes ──> Store Value
│
▼
DAWG Compression (shared suffixes)
Benefits:
- Good compression
- Supports incremental updates
- Thread-safe (atomic operations)
Build Phase:
All Keys ──> Sort ──> Build Double Array ──> Frozen Structure
Query Phase:
Key ──> Base/Check Navigation ──> Value Lookup
Benefits:
- Fastest queries
- Smallest memory footprint
- Not updatable after construction
Model ──> serde::Serialize ──> bincode ──> Bytes ──> File
Model ──> to_portable() ──> PortableModel ──> bincode ──> File
File ──> bincode ──> PortableModel ──> load_portable(factory) ──> Model<D>
The portable format stores n-grams as (key, entry) pairs, allowing reconstruction with any dictionary backend.
┌─────────────────┐
│ Sentences │
│ (Vec) │
└────────┬────────┘
│
│ par_chunks(batch_size)
▼
┌────────┬───────┬───────┬────────┐
│ Thread │Thread │Thread │ Thread │
│ 1 │ 2 │ 3 │ 4 │
└───┬────┴───┬───┴───┬───┴────┬───┘
│ │ │ │
▼ ▼ ▼ ▼
Count Count Count Count
N-grams N-grams N-grams N-grams
│ │ │ │
└────────┴───────┴────────┘
│
│ Atomic merge into dictionary
▼
┌────────────────┐
│ Dictionary │
│ (Thread-safe) │
└────────────────┘
Queries are embarrassingly parallel:
let sentences: Vec<Vec<&str>> = /* ... */;
let scores: Vec<f64> = sentences.par_iter()
.map(|s| model.sentence_log_prob(s))
.collect();word_vector(word)
│
▼
┌──────────┐
│ Cache │──hit──> Return cached vector
│ (DashMap)│
└────┬─────┘
│ miss
▼
Compute vector
│
▼
Store in cache (if space)
│
▼
Return vector
score(word, context)
│
▼
┌──────────┐
│ LRU │──hit──> Return cached score
│ Cache │
└────┬─────┘
│ miss
▼
Compute score
│
▼
Store in cache
│
▼
Return score
- Threading Model - Concurrency details
- NgramModel API - Query interface
- HybridLanguageModel API - Hybrid interface