Skip to main content

zsfm_lag_llama/
config.rs

1/// Lag-Llama model configuration (fixed from checkpoint).
2#[derive(Clone)]
3pub struct LagLlamaConfig {
4    pub n_layer: usize,
5    pub n_head: usize,
6    pub n_embd_per_head: usize,
7    pub n_embd: usize,
8    pub mlp_hidden: usize,
9    pub feature_size: usize,
10    pub n_lags: usize,
11    pub n_time_feat: usize,
12    pub max_context_length: usize,
13    pub lags_seq: Vec<usize>,
14}
15
16impl LagLlamaConfig {
17    /// Values extracted from the lag-llama.ckpt hyper_parameters.
18    pub fn default_from_ckpt() -> Self {
19        let lags_seq: Vec<usize> = vec![
20            0, 7, 8, 10, 11, 12, 13, 14, 19, 20, 21, 22, 23, 24, 26, 27, 28, 29, 30,
21            34, 35, 36, 46, 47, 48, 50, 51, 52, 55, 57, 58, 59, 60, 61, 70, 71, 72,
22            83, 94, 95, 96, 102, 103, 104, 117, 118, 119, 120, 121, 142, 143, 144,
23            154, 155, 156, 166, 167, 168, 177, 178, 179, 180, 181, 334, 335, 336,
24            362, 363, 364, 502, 503, 504, 670, 671, 672, 718, 719, 720, 726, 727,
25            728, 1090, 1091, 1092,
26        ];
27        let n_lags = lags_seq.len(); // 84
28        let feature_size = 92;
29        let n_time_feat = feature_size - n_lags; // 8
30        Self {
31            n_layer: 8,
32            n_head: 9,
33            n_embd_per_head: 16,
34            n_embd: 144,       // n_head * n_embd_per_head
35            mlp_hidden: 512,
36            feature_size,
37            n_lags,
38            n_time_feat,
39            max_context_length: 2048,
40            lags_seq,
41        }
42    }
43}