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, 34, 35, 36,
21            46, 47, 48, 50, 51, 52, 55, 57, 58, 59, 60, 61, 70, 71, 72, 83, 94, 95, 96, 102, 103,
22            104, 117, 118, 119, 120, 121, 142, 143, 144, 154, 155, 156, 166, 167, 168, 177, 178,
23            179, 180, 181, 334, 335, 336, 362, 363, 364, 502, 503, 504, 670, 671, 672, 718, 719,
24            720, 726, 727, 728, 1090, 1091, 1092,
25        ];
26        let n_lags = lags_seq.len(); // 84
27        let feature_size = 92;
28        let n_time_feat = feature_size - n_lags; // 8
29        Self {
30            n_layer: 8,
31            n_head: 9,
32            n_embd_per_head: 16,
33            n_embd: 144, // n_head * n_embd_per_head
34            mlp_hidden: 512,
35            feature_size,
36            n_lags,
37            n_time_feat,
38            max_context_length: 2048,
39            lags_seq,
40        }
41    }
42}