ferritin_plms/amplify/
config.rs

1use candle_nn::Activation;
2use serde::Deserialize;
3
4#[derive(Debug, Clone, Deserialize)]
5/// Configuration Struct for AMPLIFY
6///
7/// Currently only holds the weight params for
8/// those models found on GH: the 120M and 350M models.
9///
10pub struct AMPLIFYConfig {
11    pub hidden_size: usize,
12    pub num_hidden_layers: usize,
13    pub num_attention_heads: usize,
14    pub intermediate_size: usize,
15    pub dropout_prob: f64,
16    pub embedding_init_range: f64,
17    pub decoder_init_range: f64,
18    pub rms_norm: bool,
19    pub norm_eps: f64,
20    pub hidden_act: Activation,
21    /// Apply an extra LayerNorm after the token embedding.
22    /// Both 120M and 350M pretrained models set this to `false`; no corresponding
23    /// weight exists in their safetensors, so this field is intentionally unused
24    /// at inference time.
25    pub layer_norm_after_embedding: bool,
26    pub layer_norm_before_last_layer: bool,
27    pub vocab_size: usize,
28    /// Whether Q/K/V/output projections include a bias term.
29    pub ffn_bias: bool,
30    /// Whether FFN (w12, w3) projections include a bias term.
31    pub att_bias: bool,
32    pub pad_token_id: usize,
33    pub max_length: usize,
34}
35
36impl Default for AMPLIFYConfig {
37    fn default() -> Self {
38        AMPLIFYConfig::amp_120m()
39    }
40}
41impl AMPLIFYConfig {
42    pub fn amp_120m() -> Self {
43        Self {
44            hidden_size: 640,
45            num_hidden_layers: 24,
46            num_attention_heads: 10,
47            intermediate_size: 2560,
48            dropout_prob: 0.0,
49            embedding_init_range: 0.02,
50            decoder_init_range: 0.02,
51            rms_norm: true,
52            norm_eps: 1e-5,
53            hidden_act: Activation::Swiglu,
54            layer_norm_after_embedding: false,
55            layer_norm_before_last_layer: true,
56            vocab_size: 27,
57            ffn_bias: false,
58            att_bias: false,
59            pad_token_id: 0,
60            max_length: 2048,
61        }
62    }
63    pub fn amp_350m() -> Self {
64        Self {
65            hidden_size: 960,
66            num_hidden_layers: 32,
67            num_attention_heads: 15,
68            intermediate_size: 3840,
69            dropout_prob: 0.0,
70            embedding_init_range: 0.02,
71            decoder_init_range: 0.02,
72            rms_norm: true,
73            norm_eps: 1e-5,
74            hidden_act: Activation::Swiglu,
75            layer_norm_after_embedding: false,
76            layer_norm_before_last_layer: true,
77            vocab_size: 27,
78            ffn_bias: false,
79            att_bias: false,
80            pad_token_id: 0,
81            max_length: 2048,
82        }
83    }
84}
ferritin_plms/amplify/config.rs

ferritin_plms/amplify/
config.rs
