| 77 | #[async_trait] |
| 78 | impl Forwarder for Transformer { |
| 79 | fn load(name: String, ctx: &Context) -> Result<Box<Self>> { |
| 80 | let vb = ctx |
| 81 | .var_builder |
| 82 | .as_ref() |
| 83 | .expect("No var_builder specified") |
| 84 | .pp(&name); |
| 85 | let cfg = ctx.config.as_ref().expect("No config specified"); |
| 86 | |
| 87 | let attn = super::CausalSelfAttention::load(vb.pp("self_attn"), cfg, ctx.backend.clone())?; |
| 88 | let mlp = super::MLP::load(vb.pp("mlp"), cfg, ctx.backend.clone())?; |
| 89 | let rms_1_weight = vb.pp("input_layernorm").get(cfg.hidden_size, "weight")?; |
| 90 | let rms_2_weight = vb.pp("post_attention_layernorm").get(cfg.hidden_size, "weight")?; |
| 91 | let rms_eps = cfg.rms_norm_eps as f32; |
| 92 | Ok(Box::new(Self { |
| 93 | name, |
| 94 | rms_1_weight, |
| 95 | rms_2_weight, |
| 96 | rms_eps, |
| 97 | attn, |
| 98 | mlp, |
| 99 | backend: ctx.backend.clone(), |
| 100 | })) |
| 101 | } |
| 102 | |
| 103 | async fn forward( |
| 104 | &self, |