| 36 | #[async_trait] |
| 37 | impl Forwarder for OLMo2Block { |
| 38 | fn load(name: String, ctx: &Context) -> Result<Box<Self>> { |
| 39 | let vb = ctx |
| 40 | .var_builder |
| 41 | .as_ref() |
| 42 | .expect("No var_builder specified") |
| 43 | .pp(&name); |
| 44 | let cfg = ctx.config.as_ref().expect("No config specified"); |
| 45 | |
| 46 | let attn = CausalSelfAttention::load(vb.pp("self_attn"), cfg, ctx.backend.clone())?; |
| 47 | let mlp = MLP::load(vb.pp("mlp"), cfg, ctx.backend.clone())?; |
| 48 | |
| 49 | let h = cfg.hidden_size; |
| 50 | |
| 51 | let post_attention_layernorm_weight = |
| 52 | vb.pp("post_attention_layernorm").get(h, "weight")?; |
| 53 | let post_feedforward_layernorm_weight = |
| 54 | vb.pp("post_feedforward_layernorm").get(h, "weight")?; |
| 55 | let rms_eps = cfg.rms_norm_eps as f32; |
| 56 | |
| 57 | Ok(Box::new(Self { |
| 58 | name, |
| 59 | attn, |
| 60 | post_attention_layernorm_weight, |
| 61 | mlp, |
| 62 | post_feedforward_layernorm_weight, |
| 63 | rms_eps, |
| 64 | })) |
| 65 | } |
| 66 | |
| 67 | async fn forward( |
| 68 | &self, |