libgrammstein 0.1.0

Hybrid language model (N-gram + Embeddings) for WFST text correction
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

//! `train hybrid` command: combine n-gram and embedding models.

use crate::cli::args::TrainHybridArgs;
use crate::cli::error::{print_success, CliError, CliResult};
use crate::ngram::NgramEntry;

pub(super) fn train_hybrid(args: TrainHybridArgs, verbose: bool, quiet: bool) -> CliResult<()> {
    use crate::cli::args::InterpolationStrategy as CliStrategy;
    use crate::embedding::SubwordEmbedding;
    use crate::hybrid::{
        HybridConfig, HybridLanguageModel, InterpolationStrategy as HybridStrategy,
    };
    use crate::ngram::NgramModel;
    use liblevenshtein::dictionary::dynamic_dawg_char::DynamicDawgChar;

    if verbose {
        eprintln!("Creating hybrid model");
        eprintln!("  N-gram:    {}", args.ngram_model.display());
        eprintln!("  Embedding: {}", args.embedding_model.display());
        eprintln!("  Output:    {}", args.output.display());
        eprintln!("  Strategy:  {:?}", args.strategy);
        eprintln!("  Alpha:     {}", args.alpha);
        eprintln!("  Cache:     {}", args.cache_size);
    }

    // Check that input models exist
    if !args.ngram_model.exists() {
        return Err(CliError::file_not_found(&args.ngram_model));
    }
    if !args.embedding_model.exists() {
        return Err(CliError::file_not_found(&args.embedding_model));
    }

    if !quiet {
        eprintln!("Loading N-gram model...");
    }

    // Load N-gram model using portable format (with DynamicDawgChar backend)
    let ngram_model: NgramModel<DynamicDawgChar<NgramEntry>> =
        NgramModel::load_portable(&args.ngram_model, DynamicDawgChar::new)
            .map_err(|e| CliError::io(format!("Failed to load N-gram model: {}", e)))?;

    if verbose {
        eprintln!(
            "  Loaded N-gram model: order={}, vocab={}, ngrams={}",
            ngram_model.order(),
            ngram_model.vocab_size(),
            ngram_model.ngram_count()
        );
    }

    if !quiet {
        eprintln!("Loading embedding model...");
    }

    // Load embedding model
    let embedding_model = SubwordEmbedding::load(&args.embedding_model)
        .map_err(|e| CliError::io(format!("Failed to load embedding model: {}", e)))?;

    if verbose {
        eprintln!(
            "  Loaded embedding model: vocab={}, dim={}",
            embedding_model.vocab_size(),
            embedding_model.dim()
        );
    }

    // Map CLI interpolation strategy to hybrid module strategy
    let strategy = match args.strategy {
        CliStrategy::Linear => HybridStrategy::Linear { alpha: args.alpha },
        CliStrategy::LogLinear => HybridStrategy::LogLinear { alpha: args.alpha },
        CliStrategy::NgramFallback => HybridStrategy::NgramWithEmbeddingFallback,
        CliStrategy::Dynamic => HybridStrategy::Dynamic {
            base_alpha: args.alpha * 0.5,
            alpha_per_context: 0.1,
            max_alpha: args.alpha.min(0.95),
        },
    };

    // Create hybrid configuration
    let config = HybridConfig {
        strategy,
        cache_size: args.cache_size,
        ..Default::default()
    };

    if !quiet {
        eprintln!("Creating hybrid model...");
    }

    // Create hybrid model
    let hybrid_model = HybridLanguageModel::new(ngram_model, embedding_model, config);

    if !quiet {
        eprintln!("Saving hybrid model...");
    }

    // Save hybrid model using portable format (works with any dictionary backend)
    hybrid_model
        .save_portable(&args.output)
        .map_err(|e| CliError::io(format!("Failed to save hybrid model: {}", e)))?;

    // Get file size for display
    let file_size = std::fs::metadata(&args.output)
        .map(|m| m.len())
        .unwrap_or(0);
    let size_str = humansize::format_size(file_size, humansize::BINARY);

    if !quiet {
        print_success(&format!("Hybrid model saved to: {}", args.output.display()));
        eprintln!(
            "  N-gram order:        {}",
            hybrid_model.ngram_model().order()
        );
        eprintln!(
            "  N-gram vocab:        {}",
            hybrid_model.ngram_model().vocab_size()
        );
        eprintln!(
            "  Embedding dim:       {}",
            hybrid_model.embedding_model().dim()
        );
        eprintln!(
            "  Embedding vocab:     {}",
            hybrid_model.embedding_model().vocab_size()
        );
        eprintln!("  Strategy:            {:?}", args.strategy);
        eprintln!("  Alpha:               {}", args.alpha);
        eprintln!("  Cache size:          {}", args.cache_size);
        eprintln!("  File size:           {}", size_str);
    }

    Ok(())
}