Skip to main content

stateful_scoring/
stateful_scoring.rs

1use kenlm::Model;
2use std::env;
3
4fn main() -> Result<(), kenlm::KenlmError> {
5    let mut args = env::args().skip(1);
6    let model_path = args.next().unwrap_or_else(|| "lm/test.arpa".to_string());
7    let words = args.collect::<Vec<_>>();
8    let words = if words.is_empty() {
9        vec!["looking".to_string(), "on".to_string(), "a".to_string()]
10    } else {
11        words
12    };
13
14    let model = Model::new(model_path)?;
15    let mut state = model.begin_sentence_state();
16    let mut next = model.null_context_state();
17    let mut total = 0.0;
18
19    for word in &words {
20        let word_index = model.index(word)?;
21        let full = model.base_full_score(&state, word_index, &mut next)?;
22        total += full.log_prob;
23        println!(
24            "{word}\tindex={word_index}\tlog10={:.6}\tngram_length={}",
25            full.log_prob, full.ngram_length
26        );
27        std::mem::swap(&mut state, &mut next);
28    }
29
30    let eos = model.base_full_score(&state, model.end_sentence_index(), &mut next)?;
31    total += eos.log_prob;
32    println!(
33        "</s>\tindex={}\tlog10={:.6}\tngram_length={}",
34        model.end_sentence_index(),
35        eos.log_prob,
36        eos.ngram_length
37    );
38    println!("total: {total}");
39
40    Ok(())
41}